1use crate::ledger::MemoryStore;
8use crate::recall::{RecallQuery, RecallResult};
9use crate::trace::Trace;
10use serde::{Deserialize, Serialize};
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct FocusItem {
14 pub trace: Trace,
15 pub activation: f32,
16 pub pinned: bool,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct FocusCapsule {
21 pub capacity: usize,
22 pub items: Vec<FocusItem>,
23}
24
25impl FocusCapsule {
26 pub fn new(capacity: usize) -> Self {
27 Self {
28 capacity: capacity.max(1),
29 items: Vec::new(),
30 }
31 }
32
33 pub fn from_recall(results: Vec<RecallResult>, capacity: usize) -> Self {
34 let mut capsule = Self::new(capacity);
35 capsule.absorb(results);
36 capsule
37 }
38
39 pub fn absorb(&mut self, results: Vec<RecallResult>) {
42 for result in results {
43 if let Some(item) = self
44 .items
45 .iter_mut()
46 .find(|item| item.trace.id == result.trace.id)
47 {
48 item.activation = item.activation.max(result.score);
49 continue;
50 }
51 self.items.push(FocusItem {
52 trace: result.trace,
53 activation: result.score,
54 pinned: false,
55 });
56 }
57 self.items.sort_by(|a, b| {
58 b.pinned.cmp(&a.pinned).then_with(|| {
59 b.activation
60 .partial_cmp(&a.activation)
61 .unwrap_or(std::cmp::Ordering::Equal)
62 })
63 });
64 let pinned = self.items.iter().filter(|item| item.pinned).count();
65 self.items.truncate(self.capacity.max(pinned));
66 }
67
68 pub fn pin(&mut self, trace_id: &str) -> bool {
69 if let Some(item) = self.items.iter_mut().find(|item| item.trace.id == trace_id) {
70 item.pinned = true;
71 true
72 } else {
73 false
74 }
75 }
76
77 pub fn unpin(&mut self, trace_id: &str) -> bool {
78 if let Some(item) = self.items.iter_mut().find(|item| item.trace.id == trace_id) {
79 item.pinned = false;
80 true
81 } else {
82 false
83 }
84 }
85}
86
87impl MemoryStore {
88 pub fn focus(&self, query: &RecallQuery) -> FocusCapsule {
90 FocusCapsule::from_recall(self.recall(query), self.policy().working_set_limit)
91 }
92}