1use super::local_store::LocalKnowledgeStore;
4use super::query::KnowledgeQuery;
5use super::snapshot::KnowledgeSnapshot;
6use lean_ctx_protocol::knowledge::{ClassificationLevel, KnowledgeObjectV1};
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct RetrievalPolicy {
12 pub max_items: usize,
14 #[serde(default)]
16 pub allowed_classifications: Vec<ClassificationLevel>,
17 pub required_validity: bool,
19 pub recency_weight: f32,
21}
22
23impl Default for RetrievalPolicy {
24 fn default() -> Self {
25 Self {
26 max_items: usize::MAX,
27 allowed_classifications: Vec::new(),
28 required_validity: true,
29 recency_weight: 0.0,
30 }
31 }
32}
33
34#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
36pub struct RetrievalResult {
37 pub items: Vec<KnowledgeObjectV1>,
38 pub total_available: usize,
39 pub budget_used: usize,
40 pub excluded_reasons: Vec<String>,
41}
42
43pub trait KnowledgeRetriever {
45 fn retrieve_for_task(
47 &self,
48 task_id: &str,
49 budget: usize,
50 policy: &RetrievalPolicy,
51 ) -> Vec<KnowledgeObjectV1>;
52}
53
54impl LocalKnowledgeStore {
55 pub fn retrieve_for_task_with_result(
57 &self,
58 _task_id: &str,
59 budget: usize,
60 policy: &RetrievalPolicy,
61 ) -> RetrievalResult {
62 let now_query = if policy.required_validity {
63 Some(KnowledgeQuery::valid_at(chrono::Utc::now()))
64 } else {
65 None
66 };
67 let cap = budget.min(policy.max_items);
68 let mut excluded_reasons = Vec::new();
69 let mut eligible = Vec::new();
70
71 for object in self.values() {
72 if let Some(query) = &now_query {
73 if !query.matches(object) {
74 excluded_reasons.push(format!("{}:invalid_or_superseded", object.object_id()));
75 continue;
76 }
77 }
78 if !policy.allowed_classifications.is_empty()
79 && object.classification.as_ref().is_none_or(|classification| {
80 !policy
81 .allowed_classifications
82 .contains(&classification.level)
83 })
84 {
85 excluded_reasons.push(format!("{}:classification", object.object_id()));
86 continue;
87 }
88 eligible.push(object.clone());
89 }
90
91 let total_available = eligible.len();
92 if eligible.len() > cap {
93 excluded_reasons.push(format!("budget:{cap}"));
94 eligible.truncate(cap);
95 }
96
97 let _recency_weight = policy.recency_weight;
100 RetrievalResult {
101 budget_used: eligible.len(),
102 items: eligible,
103 total_available,
104 excluded_reasons,
105 }
106 }
107
108 pub fn snapshot_for_task(
110 &self,
111 task_id: &str,
112 budget: usize,
113 policy: &RetrievalPolicy,
114 policy_version: &str,
115 ) -> KnowledgeSnapshot {
116 let result = self.retrieve_for_task_with_result(task_id, budget, policy);
117 KnowledgeSnapshot::from_items(task_id, policy_version, &result.items)
118 }
119}
120
121impl KnowledgeRetriever for LocalKnowledgeStore {
122 fn retrieve_for_task(
123 &self,
124 task_id: &str,
125 budget: usize,
126 policy: &RetrievalPolicy,
127 ) -> Vec<KnowledgeObjectV1> {
128 self.retrieve_for_task_with_result(task_id, budget, policy)
129 .items
130 }
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136 use crate::core::knowledge::store::KnowledgeStore;
137 use chrono::{Duration, Utc};
138 use lean_ctx_protocol::knowledge::{
139 ClassificationLevel, DataClassification, KnowledgeSourceType, SourceReference,
140 ValidityWindow,
141 };
142
143 fn object(id: &str, level: ClassificationLevel, valid: bool) -> KnowledgeObjectV1 {
144 let now = Utc::now();
145 KnowledgeObjectV1 {
146 schema_version: 1,
147 source_ref: Some(SourceReference {
148 source_type: "test".to_owned(),
149 uri: "memory://knowledge".to_owned(),
150 commit_sha: None,
151 access_timestamp: now.to_rfc3339(),
152 }),
153 source_type: KnowledgeSourceType::Other,
154 authority: None,
155 owner: "test".to_owned(),
156 classification: Some(DataClassification {
157 level,
158 retention_days: None,
159 }),
160 validity: valid.then(|| ValidityWindow {
161 valid_from: (now - Duration::minutes(1)).to_rfc3339(),
162 valid_until: Some((now + Duration::minutes(1)).to_rfc3339()),
163 superseded_by: None,
164 }),
165 supersedes: None,
166 content_hash: id.to_owned(),
167 evidence_digest: format!("digest:{id}"),
168 policy_ref: "policy:test".to_owned(),
169 evidence_refs: Vec::new(),
170 extra: Default::default(),
171 }
172 }
173
174 #[test]
175 fn local_retrieval_applies_policy_and_budget() {
176 let mut store = LocalKnowledgeStore::new();
177 store
178 .put(object("public", ClassificationLevel::Public, true))
179 .expect("put public");
180 store
181 .put(object("internal", ClassificationLevel::Internal, true))
182 .expect("put internal");
183 store
184 .put(object("expired", ClassificationLevel::Public, false))
185 .expect("put expired");
186
187 let policy = RetrievalPolicy {
188 max_items: 1,
189 allowed_classifications: vec![ClassificationLevel::Public],
190 required_validity: true,
191 recency_weight: 0.5,
192 };
193 let result = store.retrieve_for_task_with_result("task-1", 4, &policy);
194 assert_eq!(result.items.len(), 1);
195 assert_eq!(result.items[0].object_id(), "public");
196 assert_eq!(result.total_available, 1);
197 assert_eq!(result.budget_used, 1);
198 assert!(
199 result
200 .excluded_reasons
201 .iter()
202 .any(|reason| reason.contains("invalid"))
203 );
204 }
205}