1#![allow(missing_docs)]
2use 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
27pub 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 pub fn execute_with_time(
57 &self,
58 req: QuerySubjectRequest,
59 now: DateTime<Utc>,
60 ) -> Result<QuerySubjectResponse, MemError> {
61 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 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", _ => "NoBelief",
93 };
94
95 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 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 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 entries.sort_by(|a, b| a.predicate.cmp(&b.predicate));
147
148 Ok(QuerySubjectResponse { entries })
149 }
150
151 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#[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 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]
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 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 assert_eq!(preds, vec!["city", "dietary", "employer"], "entries must be sorted by predicate");
302 }
303
304 #[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(); 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(); 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]
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 let tx = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
347
348 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 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 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]
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 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 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}