Skip to main content

lean_ctx/core/knowledge/
retrieval.rs

1//! Retrieval contract and the local reference implementation.
2
3use 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/// Policy inputs understood by the portable retrieval surface.
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct RetrievalPolicy {
12    /// Maximum number of objects returned by a retrieval.
13    pub max_items: usize,
14    /// Empty means every classification is allowed.
15    #[serde(default)]
16    pub allowed_classifications: Vec<ClassificationLevel>,
17    /// If true, superseded and out-of-window objects are excluded.
18    pub required_validity: bool,
19    /// Reserved portable ranking hint; enterprise implementations may use it.
20    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/// Detailed retrieval accounting suitable for a task-local snapshot.
35#[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
43/// Portable retrieval interface implemented by local and enterprise adapters.
44pub trait KnowledgeRetriever {
45    /// Retrieve at most `budget` valid objects for a task.
46    fn retrieve_for_task(
47        &self,
48        task_id: &str,
49        budget: usize,
50        policy: &RetrievalPolicy,
51    ) -> Vec<KnowledgeObjectV1>;
52}
53
54impl LocalKnowledgeStore {
55    /// Retrieve with accounting details from the local reference store.
56    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        // The local reference implementation deliberately keeps stable store order;
98        // `recency_weight` is available to richer adapters without changing this MVP.
99        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    /// Produce a receipt-ready proof of the objects selected for a task.
109    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}