Skip to main content

rig_resources/
projection.rs

1//! Projection helpers for `rig-compose` context packing.
2
3use rig_compose::{
4    ContextItem, ContextPack, ContextPackConfig, ContextSourceKind, Evidence, InvestigationContext,
5};
6use serde_json::{Value, json};
7
8use crate::{BehaviorPattern, EntityBaseline, MemoryLookupHit};
9
10/// Convert resource-native records into [`ContextItem`] values.
11pub trait IntoContextItem {
12    /// Project this resource record into a prompt-ready context item.
13    fn to_context_item(&self) -> ContextItem;
14}
15
16impl IntoContextItem for BehaviorPattern {
17    fn to_context_item(&self) -> ContextItem {
18        let source_id = format!("behavior_pattern/{}@v{}", self.id, self.version);
19        let text = if self.description.is_empty() {
20            format!("behavior pattern {} version {}", self.id, self.version)
21        } else {
22            self.description.clone()
23        };
24        ContextItem::new(ContextSourceKind::Resource, source_id, text)
25            .with_score(f64::from(self.confidence_delta))
26            .with_provenance(json!({
27                "resource": "behavior_pattern",
28                "id": self.id,
29                "version": self.version,
30                "required": self.rule.required,
31                "forbidden": self.rule.forbidden,
32                "confidence_delta": self.confidence_delta,
33                "conclude": self.conclude,
34            }))
35    }
36}
37
38impl IntoContextItem for EntityBaseline {
39    fn to_context_item(&self) -> ContextItem {
40        ContextItem::new(
41            ContextSourceKind::Resource,
42            format!("baseline/{}/{}", self.entity, self.metric),
43            format!(
44                "baseline for {} {}: mean {}, std_dev {}, samples {}",
45                self.entity, self.metric, self.mean, self.std_dev, self.samples
46            ),
47        )
48        .with_score(self.samples as f64)
49        .with_provenance(json!({
50            "resource": "baseline",
51            "entity": self.entity,
52            "metric": self.metric,
53            "mean": self.mean,
54            "std_dev": self.std_dev,
55            "samples": self.samples,
56        }))
57    }
58}
59
60impl IntoContextItem for MemoryLookupHit {
61    fn to_context_item(&self) -> ContextItem {
62        memory_hit_to_context_item(self, 0)
63    }
64}
65
66/// Project a memory lookup hit into a ranked memory context item.
67#[must_use]
68pub fn memory_hit_to_context_item(hit: &MemoryLookupHit, rank: usize) -> ContextItem {
69    let source_id = hit
70        .key
71        .clone()
72        .unwrap_or_else(|| format!("memory.hit/{rank}"));
73    ContextItem::new(ContextSourceKind::Memory, source_id, hit.summary.clone())
74        .with_rank(rank)
75        .with_score(f64::from(hit.score))
76        .with_provenance(json!({
77            "resource": "memory.lookup",
78            "key": hit.key,
79            "score": hit.score,
80            "metadata": hit.metadata,
81        }))
82}
83
84/// Project memory lookup hits into ranked memory context items.
85#[must_use]
86pub fn memory_hits_to_context_items(hits: &[MemoryLookupHit]) -> Vec<ContextItem> {
87    hits.iter()
88        .enumerate()
89        .map(|(rank, hit)| memory_hit_to_context_item(hit, rank))
90        .collect()
91}
92
93/// Project all accumulated investigation evidence into resource or memory
94/// context items.
95#[must_use]
96pub fn evidence_to_context_items(ctx: &InvestigationContext) -> Vec<ContextItem> {
97    ctx.evidence
98        .iter()
99        .enumerate()
100        .map(|(rank, evidence)| evidence_to_context_item(evidence, rank))
101        .collect()
102}
103
104/// Project one evidence record into a context item.
105#[must_use]
106pub fn evidence_to_context_item(evidence: &Evidence, rank: usize) -> ContextItem {
107    let source = if evidence.source_skill == "general.memory_pivot"
108        || evidence.label.starts_with("memory.")
109    {
110        ContextSourceKind::Memory
111    } else {
112        ContextSourceKind::Resource
113    };
114    let source_id = format!("{}/{}", evidence.source_skill, evidence.label);
115    ContextItem::new(source, source_id, evidence_text(evidence))
116        .with_rank(rank)
117        .with_score(evidence_score(&evidence.detail))
118        .with_provenance(json!({
119            "source_skill": evidence.source_skill,
120            "label": evidence.label,
121            "detail": evidence.detail,
122        }))
123}
124
125/// Pack resource-projected context items with the shared kernel packer.
126#[must_use]
127pub fn pack_resource_context(items: Vec<ContextItem>, config: ContextPackConfig) -> ContextPack {
128    ContextPack::pack(items, config)
129}
130
131fn evidence_text(evidence: &Evidence) -> String {
132    evidence
133        .detail
134        .get("summary")
135        .and_then(Value::as_str)
136        .or_else(|| evidence.detail.get("description").and_then(Value::as_str))
137        .map(str::to_owned)
138        .unwrap_or_else(|| evidence.label.clone())
139}
140
141fn evidence_score(detail: &Value) -> f64 {
142    detail
143        .get("score")
144        .and_then(Value::as_f64)
145        .or_else(|| detail.get("delta").and_then(Value::as_f64))
146        .or_else(|| detail.get("confidence_delta").and_then(Value::as_f64))
147        .unwrap_or(0.0)
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::PatternRule;
154    use rig_compose::ContextOmissionReason;
155
156    #[test]
157    fn behavior_pattern_projects_to_resource_context() {
158        let pattern = BehaviorPattern::new(
159            "spray",
160            2,
161            PatternRule {
162                required: vec!["auth.failure.burst".into()],
163                forbidden: vec!["baseline.within".into()],
164            },
165            0.25,
166        )
167        .with_description("password spray around one host");
168
169        let item = pattern.to_context_item();
170
171        assert_eq!(item.source, ContextSourceKind::Resource);
172        assert_eq!(item.source_id, "behavior_pattern/spray@v2");
173        assert_eq!(item.text, "password spray around one host");
174        assert!((item.score - 0.25).abs() < 1e-9);
175        assert_eq!(item.provenance["resource"], "behavior_pattern");
176        assert_eq!(item.provenance["required"][0], "auth.failure.burst");
177    }
178
179    #[test]
180    fn memory_hits_project_with_stable_ranks() {
181        let hits = vec![
182            MemoryLookupHit::new(0.9, "first").with_key("episode-1"),
183            MemoryLookupHit::new(0.5, "second"),
184        ];
185
186        let items = memory_hits_to_context_items(&hits);
187
188        assert_eq!(items.len(), 2);
189        assert_eq!(items[0].source, ContextSourceKind::Memory);
190        assert_eq!(items[0].source_id, "episode-1");
191        assert_eq!(items[0].rank, 0);
192        assert_eq!(items[1].source_id, "memory.hit/1");
193        assert_eq!(items[1].rank, 1);
194    }
195
196    #[test]
197    fn evidence_projection_packs_and_omits_by_kernel_rules() {
198        let mut ctx = InvestigationContext::new("host", "partition");
199        ctx.evidence.push(
200            Evidence::new("general.memory_pivot", "memory.hit")
201                .with_detail(json!({"summary": "matching episode", "score": 0.8})),
202        );
203        ctx.evidence.push(
204            Evidence::new("knowledge.behavior_pattern", "pattern:spray")
205                .with_detail(json!({"description": "spray pattern", "delta": 0.2})),
206        );
207
208        let items = evidence_to_context_items(&ctx);
209        let pack = pack_resource_context(items, ContextPackConfig::new(1_000).with_max_items(1));
210
211        assert_eq!(pack.selected.len(), 1);
212        assert_eq!(pack.omitted.len(), 1);
213        assert_eq!(pack.omitted[0].reason, ContextOmissionReason::MaxItems);
214        assert_eq!(pack.selected[0].source, ContextSourceKind::Memory);
215        assert_eq!(pack.selected[0].text, "matching episode");
216    }
217}