Skip to main content

mempill_core/application/
query_subject.rs

1#![allow(missing_docs)]
2//! QuerySubjectUseCase — enumerate all resolved beliefs for a subject across predicates.
3//!
4//! Read-only. For each distinct predicate stored under the subject (filtered by the
5//! tx-time cutoff), delegates to the EXISTING query_memory fold (QueryMemoryUseCase) and
6//! collects the per-predicate result into a SubjectFactEntry.
7//!
8//! No new fold logic — 100% reuse of the existing truth_engine fold path.
9
10use std::sync::Arc;
11
12use chrono::{DateTime, Utc};
13
14use mempill_types::{BeliefStatus, ProvenanceLabel, ExternalKind};
15
16use crate::{
17    config::EngineConfig,
18    error::MemError,
19    ports::{PersistencePort, VectorPort},
20};
21
22use super::dto::{
23    QueryMemoryRequest, QuerySubjectRequest, QuerySubjectResponse, SubjectFactEntry,
24};
25use super::query_memory::QueryMemoryUseCase;
26
27/// Use-case: return the resolved belief for every predicate under a subject.
28///
29/// Generic over persistence and vector ports (mirrors QueryMemoryUseCase).
30pub struct QuerySubjectUseCase<P, V>
31where
32    P: PersistencePort + Send + Sync + 'static,
33    V: VectorPort + Send + Sync + 'static,
34{
35    persistence: Arc<P>,
36    vector: Option<Arc<V>>,
37    config: EngineConfig,
38}
39
40impl<P, V> QuerySubjectUseCase<P, V>
41where
42    P: PersistencePort + Send + Sync + 'static,
43    V: VectorPort + Send + Sync + 'static,
44{
45    pub fn new(persistence: Arc<P>, vector: Option<Arc<V>>, config: EngineConfig) -> Self {
46        Self { persistence, vector, config }
47    }
48
49    /// Execute with an explicit `now` (DETERMINISM — no clock reads here).
50    ///
51    /// Algorithm:
52    ///   1. list_predicates_for_subject → distinct predicates visible at as_of_tx_time.
53    ///   2. For each predicate, run the existing QueryMemoryUseCase fold (same valid_at +
54    ///      as_of_tx_time), collecting a SubjectFactEntry.
55    ///   3. Sort entries by predicate (stable, deterministic output).
56    pub fn execute_with_time(
57        &self,
58        req: QuerySubjectRequest,
59        now: DateTime<Utc>,
60    ) -> Result<QuerySubjectResponse, MemError> {
61        // Step 1: list distinct predicates visible at the tx-time cutoff.
62        let predicates = self.persistence
63            .list_predicates_for_subject(&req.agent_id, &req.subject, req.as_of_tx_time)
64            .map_err(|e| MemError::Persistence { source: Box::new(e) })?;
65
66        // Step 2: fold each predicate using the existing QueryMemoryUseCase.
67        let query_uc = QueryMemoryUseCase::new(
68            Arc::clone(&self.persistence),
69            self.vector.clone(),
70            self.config.clone(),
71        );
72
73        let mut entries: Vec<SubjectFactEntry> = predicates
74            .into_iter()
75            .map(|predicate| {
76                let mem_req = QueryMemoryRequest {
77                    agent_id: req.agent_id.clone(),
78                    subject: req.subject.clone(),
79                    predicate: predicate.clone(),
80                    as_of_tx_time: req.as_of_tx_time,
81                    valid_at: req.valid_at,
82                };
83                let resp = query_uc.execute_with_time(mem_req, now)?;
84                let belief = resp.belief;
85
86                let status_str = match belief.status {
87                    BeliefStatus::Resolved => "Resolved",
88                    BeliefStatus::Contested => "Contested",
89                    BeliefStatus::NoBelief => "NoBelief",
90                    BeliefStatus::TimingUncertain => "TimingUncertain",
91                    BeliefStatus::Conflict => "Contested", // surface as Contested to caller
92                    _ => "NoBelief",
93                };
94
95                // Extract from the primary belief slot (present for Resolved / TimingUncertain).
96                let (value, valid_from_display, valid_until_display, provenance_str, claim_ref_str, conf) =
97                    if let Some(primary) = &belief.primary {
98                        let val_str = match &primary.fact.value {
99                            serde_json::Value::String(s) => Some(s.clone()),
100                            other => Some(other.to_string()),
101                        };
102
103                        // Compute display strings using the same helper as enrich_query_memory.
104                        let from_disp = mempill_types::time::format_valid_time_endpoint(
105                            primary.valid_time.start,
106                            primary.valid_time.start_granularity,
107                        );
108                        let until_disp = mempill_types::time::format_valid_time_endpoint(
109                            primary.valid_time.end,
110                            primary.valid_time.end_granularity,
111                        );
112
113                        let prov = provenance_label_to_str(&primary.provenance);
114                        let cr = primary.claim_ref.0.to_string();
115                        let conf_val = primary.confidence.value_confidence;
116
117                        (val_str, from_disp, until_disp, prov, Some(cr), Some(conf_val))
118                    } else {
119                        // NoBelief — pick from alternatives if Contested.
120                        let prov = belief.alternatives.first()
121                            .map(|a| provenance_label_to_str(&a.provenance))
122                            .unwrap_or_else(|| "none".to_string());
123                        let cr = belief.alternatives.first().map(|a| a.claim_ref.0.to_string());
124                        let conf_val = belief.alternatives.first().map(|a| a.confidence.value_confidence);
125                        let val_str = belief.alternatives.first().map(|a| match &a.fact.value {
126                            serde_json::Value::String(s) => s.clone(),
127                            other => other.to_string(),
128                        });
129                        (val_str, None, None, prov, cr, conf_val)
130                    };
131
132                Ok(SubjectFactEntry {
133                    predicate,
134                    value,
135                    status: status_str.to_string(),
136                    valid_from_display,
137                    valid_until_display,
138                    provenance: provenance_str,
139                    claim_ref: claim_ref_str,
140                    conf,
141                })
142            })
143            .collect::<Result<Vec<_>, MemError>>()?;
144
145        // Step 3: sort by predicate for stable, deterministic output.
146        entries.sort_by(|a, b| a.predicate.cmp(&b.predicate));
147
148        Ok(QuerySubjectResponse { entries })
149    }
150
151    /// Convenience wrapper that stamps now internally.
152    pub fn execute(&self, req: QuerySubjectRequest) -> Result<QuerySubjectResponse, MemError> {
153        self.execute_with_time(req, Utc::now())
154    }
155}
156
157fn provenance_label_to_str(p: &ProvenanceLabel) -> String {
158    match p {
159        ProvenanceLabel::ModelDerived => "ModelDerived".to_string(),
160        ProvenanceLabel::RecallReEntry => "RecallReEntry".to_string(),
161        ProvenanceLabel::External(ExternalKind::UserAsserted) => "External/UserAsserted".to_string(),
162        ProvenanceLabel::External(ExternalKind::ExternalFirstHand) => "External/ExternalFirstHand".to_string(),
163        _ => "Unknown".to_string(),
164    }
165}
166
167// ── Tests ──────────────────────────────────────────────────────────────────────
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172    use crate::noop::NoOpVector;
173    use crate::ports::persistence::Txn;
174    use chrono::TimeZone;
175    use mempill_types::{
176        AgentId, Cardinality, Claim, ClaimEdge, ClaimRef, Confidence, Criticality,
177        ExternalAnchor, ExternalKind, Fact, LedgerEntry, ProvenanceLabel, TransactionTime,
178        ValidTime, ValidityAssertion,
179    };
180    use std::sync::Mutex;
181
182    // ── MockTxn + FullMockStore ────────────────────────────────────────────────
183
184    struct MockTxn(AgentId);
185    impl Txn for MockTxn {
186        fn agent_id(&self) -> &AgentId { &self.0 }
187    }
188
189    #[derive(Debug, thiserror::Error)]
190    #[error("mock")]
191    struct MockErr;
192
193    #[derive(Default)]
194    struct FullMockStore {
195        claims: Mutex<Vec<Claim>>,
196        ledger: Mutex<Vec<LedgerEntry>>,
197    }
198
199    impl PersistencePort for FullMockStore {
200        type Transaction = MockTxn;
201        type Error = MockErr;
202
203        fn begin_atomic(&self, aid: &AgentId) -> Result<MockTxn, MockErr> { Ok(MockTxn(aid.clone())) }
204        fn append_claim(&self, _t: &mut MockTxn, c: &Claim) -> Result<ClaimRef, MockErr> {
205            self.claims.lock().unwrap().push(c.clone());
206            Ok(c.claim_ref().clone())
207        }
208        fn append_validity_assertion(&self, _t: &mut MockTxn, _a: &ValidityAssertion) -> Result<(), MockErr> { Ok(()) }
209        fn append_ledger_entry(&self, _t: &mut MockTxn, e: &LedgerEntry) -> Result<(), MockErr> {
210            self.ledger.lock().unwrap().push(e.clone()); Ok(())
211        }
212        fn append_claim_edge(&self, _t: &mut MockTxn, _e: &ClaimEdge) -> Result<(), MockErr> { Ok(()) }
213        fn commit(&self, _t: MockTxn) -> Result<(), MockErr> { Ok(()) }
214        fn rollback(&self, _t: MockTxn) -> Result<(), MockErr> { Ok(()) }
215        fn load_subject_line(&self, _aid: &AgentId, subject: &str, predicate: &str, as_of: Option<DateTime<Utc>>) -> Result<Vec<Claim>, MockErr> {
216            let claims = self.claims.lock().unwrap();
217            Ok(claims.iter()
218                .filter(|c| c.fact().subject == subject && c.fact().predicate == predicate)
219                .filter(|c| as_of.is_none_or(|t| c.transaction_time().0 <= t))
220                .cloned()
221                .collect())
222        }
223        fn load_claim(&self, _aid: &AgentId, _r: &ClaimRef) -> Result<Option<Claim>, MockErr> { Ok(None) }
224        fn load_validity_assertions_for(&self, _aid: &AgentId, _r: &ClaimRef) -> Result<Vec<ValidityAssertion>, MockErr> { Ok(vec![]) }
225        fn load_ledger(&self, _aid: &AgentId, _from: Option<&TransactionTime>, _lim: usize) -> Result<Vec<LedgerEntry>, MockErr> { Ok(vec![]) }
226        fn load_ledger_for_claims(&self, _aid: &AgentId, refs: &[ClaimRef], as_of: Option<DateTime<Utc>>) -> Result<Vec<LedgerEntry>, MockErr> {
227            let ledger = self.ledger.lock().unwrap();
228            Ok(ledger.iter()
229                .filter(|e| refs.contains(&e.claim_ref))
230                .filter(|e| as_of.is_none_or(|t| e.recorded_at.0 <= t))
231                .cloned()
232                .collect())
233        }
234        fn load_edges_for(&self, _aid: &AgentId, _r: &ClaimRef) -> Result<Vec<ClaimEdge>, MockErr> { Ok(vec![]) }
235        fn load_injected_claims(&self, _aid: &AgentId) -> Result<Vec<ClaimRef>, MockErr> { Ok(vec![]) }
236        fn load_lineage(&self, _aid: &AgentId, _r: &ClaimRef) -> Result<Vec<ClaimEdge>, MockErr> { Ok(vec![]) }
237        fn list_predicates_for_subject(&self, _aid: &AgentId, subject: &str, as_of: Option<DateTime<Utc>>) -> Result<Vec<String>, MockErr> {
238            let claims = self.claims.lock().unwrap();
239            let mut predicates: Vec<String> = claims.iter()
240                .filter(|c| c.fact().subject == subject)
241                .filter(|c| as_of.is_none_or(|t| c.transaction_time().0 <= t))
242                .map(|c| c.fact().predicate.clone())
243                .collect::<std::collections::HashSet<_>>()
244                .into_iter()
245                .collect();
246            predicates.sort();
247            Ok(predicates)
248        }
249    }
250
251    fn make_claim_at(subject: &str, predicate: &str, value: serde_json::Value, tx: DateTime<Utc>) -> Claim {
252        Claim::new(
253            ClaimRef::new_random(),
254            AgentId("agent".into()),
255            Fact { subject: subject.into(), predicate: predicate.into(), value },
256            Cardinality::Functional,
257            ProvenanceLabel::External(ExternalKind::UserAsserted),
258            ExternalAnchor { nearest_external_anchor: None, derivation_depth: 0 },
259            TransactionTime(tx),
260            ValidTime { start: None, end: None, valid_time_confidence: 0.0, start_granularity: None, end_granularity: None },
261            Confidence { value_confidence: 0.9, valid_time_confidence: 0.0 },
262            Criticality::Medium,
263            vec![],
264            None,
265            None,
266        )
267    }
268
269    // ── Test: query_subject returns all 3 predicates for alice-chen ────────────
270    #[test]
271    fn query_subject_returns_all_predicates() {
272        use std::sync::Arc;
273        let store = Arc::new(FullMockStore::default());
274        let agent = AgentId("agent".into());
275        let tx = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
276
277        // Seed 3 predicates for alice-chen
278        store.claims.lock().unwrap().extend([
279            make_claim_at("alice-chen", "city", serde_json::json!("Berlin"), tx),
280            make_claim_at("alice-chen", "employer", serde_json::json!("Acme Corp"), tx),
281            make_claim_at("alice-chen", "dietary", serde_json::json!("vegan"), tx),
282        ]);
283
284        let uc = QuerySubjectUseCase::new(Arc::clone(&store), None::<Arc<NoOpVector>>, EngineConfig::default());
285        let now = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
286        let req = QuerySubjectRequest {
287            agent_id: agent,
288            subject: "alice-chen".into(),
289            valid_at: None,
290            as_of_tx_time: None,
291        };
292        let resp = uc.execute_with_time(req, now).unwrap();
293
294        assert_eq!(resp.entries.len(), 3, "expected 3 predicates, got {}", resp.entries.len());
295        let preds: Vec<&str> = resp.entries.iter().map(|e| e.predicate.as_str()).collect();
296        assert!(preds.contains(&"city"), "must contain city");
297        assert!(preds.contains(&"employer"), "must contain employer");
298        assert!(preds.contains(&"dietary"), "must contain dietary");
299
300        // Sorted by predicate
301        assert_eq!(preds, vec!["city", "dietary", "employer"], "entries must be sorted by predicate");
302    }
303
304    // ── Test: predicate ingested AFTER as_of_tx_time is EXCLUDED ──────────────
305    #[test]
306    fn query_subject_excludes_predicate_after_tx_cutoff() {
307        use std::sync::Arc;
308        let store = Arc::new(FullMockStore::default());
309        let agent = AgentId("agent".into());
310
311        let t1 = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
312        let t2 = Utc.with_ymd_and_hms(2025, 6, 1, 0, 0, 0).unwrap(); // after cutoff
313
314        // city is visible at t1; phone is only added at t2 (after cutoff)
315        store.claims.lock().unwrap().extend([
316            make_claim_at("alice-chen", "city", serde_json::json!("Paris"), t1),
317            make_claim_at("alice-chen", "phone", serde_json::json!("+1234"), t2),
318        ]);
319
320        let uc = QuerySubjectUseCase::new(Arc::clone(&store), None::<Arc<NoOpVector>>, EngineConfig::default());
321        let now = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
322        let cutoff = Utc.with_ymd_and_hms(2025, 3, 1, 0, 0, 0).unwrap(); // between t1 and t2
323
324        let req = QuerySubjectRequest {
325            agent_id: agent,
326            subject: "alice-chen".into(),
327            valid_at: None,
328            as_of_tx_time: Some(cutoff),
329        };
330        let resp = uc.execute_with_time(req, now).unwrap();
331
332        let preds: Vec<&str> = resp.entries.iter().map(|e| e.predicate.as_str()).collect();
333        assert!(preds.contains(&"city"), "city must be present (before cutoff)");
334        assert!(!preds.contains(&"phone"), "phone must be EXCLUDED (after cutoff)");
335        assert_eq!(resp.entries.len(), 1);
336    }
337
338    // ── Test: valid_at narrows to the historical value for a predicate ─────────
339    #[test]
340    fn query_subject_valid_at_selects_historical_value() {
341        use std::sync::Arc;
342        let store = Arc::new(FullMockStore::default());
343        let agent = AgentId("agent".into());
344
345        // Alice lived in Berlin [2020,2023), then Paris [2023,∞).
346        let tx = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
347
348        // Two claims with valid-time windows; use make_claim but set valid_time manually.
349        let berlin = Claim::new(
350            ClaimRef::new_random(),
351            agent.clone(),
352            Fact { subject: "alice-chen".into(), predicate: "city".into(), value: serde_json::json!("Berlin") },
353            Cardinality::Functional,
354            ProvenanceLabel::External(ExternalKind::UserAsserted),
355            ExternalAnchor { nearest_external_anchor: None, derivation_depth: 0 },
356            TransactionTime(tx),
357            ValidTime {
358                start: Some(Utc.with_ymd_and_hms(2020, 1, 1, 0, 0, 0).unwrap()),
359                end: Some(Utc.with_ymd_and_hms(2023, 1, 1, 0, 0, 0).unwrap()),
360                valid_time_confidence: 0.9,
361                start_granularity: None,
362                end_granularity: None,
363            },
364            Confidence { value_confidence: 0.9, valid_time_confidence: 0.9 },
365            Criticality::Medium,
366            vec![],
367            None,
368            None,
369        );
370        let paris = Claim::new(
371            ClaimRef::new_random(),
372            agent.clone(),
373            Fact { subject: "alice-chen".into(), predicate: "city".into(), value: serde_json::json!("Paris") },
374            Cardinality::Functional,
375            ProvenanceLabel::External(ExternalKind::UserAsserted),
376            ExternalAnchor { nearest_external_anchor: None, derivation_depth: 0 },
377            TransactionTime(tx),
378            ValidTime {
379                start: Some(Utc.with_ymd_and_hms(2023, 1, 1, 0, 0, 0).unwrap()),
380                end: None,
381                valid_time_confidence: 0.9,
382                start_granularity: None,
383                end_granularity: None,
384            },
385            Confidence { value_confidence: 0.9, valid_time_confidence: 0.9 },
386            Criticality::Medium,
387            vec![],
388            None,
389            None,
390        );
391        store.claims.lock().unwrap().extend([berlin, paris]);
392
393        let uc = QuerySubjectUseCase::new(Arc::clone(&store), None::<Arc<NoOpVector>>, EngineConfig::default());
394        let now = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
395
396        // valid_at = 2021 → should resolve to Berlin
397        let req = QuerySubjectRequest {
398            agent_id: agent.clone(),
399            subject: "alice-chen".into(),
400            valid_at: Some(Utc.with_ymd_and_hms(2021, 6, 1, 0, 0, 0).unwrap()),
401            as_of_tx_time: None,
402        };
403        let resp = uc.execute_with_time(req, now).unwrap();
404        let city_entry = resp.entries.iter().find(|e| e.predicate == "city").unwrap();
405        assert_eq!(city_entry.value.as_deref(), Some("Berlin"), "valid_at=2021 must resolve to Berlin");
406        assert_eq!(city_entry.status, "Resolved");
407
408        // valid_at = 2024 → should resolve to Paris
409        let req2 = QuerySubjectRequest {
410            agent_id: agent,
411            subject: "alice-chen".into(),
412            valid_at: Some(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap()),
413            as_of_tx_time: None,
414        };
415        let resp2 = uc.execute_with_time(req2, now).unwrap();
416        let city_entry2 = resp2.entries.iter().find(|e| e.predicate == "city").unwrap();
417        assert_eq!(city_entry2.value.as_deref(), Some("Paris"), "valid_at=2024 must resolve to Paris");
418    }
419
420    // ── Test: Contested predicate reports status=Contested ────────────────────
421    #[test]
422    fn query_subject_contested_predicate_reports_contested() {
423        use std::sync::Arc;
424        use mempill_types::{Disposition, LedgerEventKind};
425
426        let store = Arc::new(FullMockStore::default());
427        let agent = AgentId("agent".into());
428        let tx = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
429
430        // Two overlapping (Functional, no valid_time) claims for "employer" — contested scenario.
431        let c1 = make_claim_at("alice-chen", "employer", serde_json::json!("Acme"), tx);
432        let c2 = make_claim_at("alice-chen", "employer", serde_json::json!("GlobalCorp"), tx);
433        let ref1 = c1.claim_ref().clone();
434        let ref2 = c2.claim_ref().clone();
435        store.claims.lock().unwrap().extend([c1, c2]);
436
437        // Mark both as Contested in the ledger so the fold surfaces Contested.
438        let le1 = LedgerEntry {
439            entry_id: uuid::Uuid::new_v4(),
440            agent_id: agent.clone(),
441            claim_ref: ref1,
442            event_kind: LedgerEventKind::ClaimCommitted,
443            disposition: Disposition::Contested,
444            rationale: None,
445            recorded_at: TransactionTime(tx),
446        };
447        let le2 = LedgerEntry {
448            entry_id: uuid::Uuid::new_v4(),
449            agent_id: agent.clone(),
450            claim_ref: ref2,
451            event_kind: LedgerEventKind::ClaimCommitted,
452            disposition: Disposition::Contested,
453            rationale: None,
454            recorded_at: TransactionTime(tx),
455        };
456        store.ledger.lock().unwrap().extend([le1, le2]);
457
458        let uc = QuerySubjectUseCase::new(Arc::clone(&store), None::<Arc<NoOpVector>>, EngineConfig::default());
459        let now = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
460        let req = QuerySubjectRequest {
461            agent_id: agent,
462            subject: "alice-chen".into(),
463            valid_at: None,
464            as_of_tx_time: None,
465        };
466        let resp = uc.execute_with_time(req, now).unwrap();
467        let employer_entry = resp.entries.iter().find(|e| e.predicate == "employer").unwrap();
468        assert_eq!(employer_entry.status, "Contested", "employer must be Contested");
469    }
470}