Skip to main content

mempal_runtime/ingest/
diary.rs

1use crate::core::{
2    db::Database,
3    types::{AnchorKind, Drawer, MemoryDomain, Provenance, SourceType},
4    utils::current_timestamp,
5};
6use crate::embed::Embedder;
7use crate::ingest::lock;
8use crate::ingest::normalize::CURRENT_NORMALIZE_VERSION;
9use crate::ingest::{IngestError, IngestStats, LOCK_TIMEOUT, mempal_home_from_db};
10
11pub const DIARY_ROLLUP_WING: &str = "agent-diary";
12pub const DAILY_ROLLUP_LIMIT_BYTES: usize = 32 * 1024;
13const DIARY_ROLLUP_SEPARATOR: &str = "\n\n---\n\n";
14
15#[derive(Debug, Clone, Copy)]
16pub struct DiaryRollupOptions<'a> {
17    pub room: Option<&'a str>,
18    pub day: Option<&'a str>,
19    pub dry_run: bool,
20    pub importance: i32,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct DiaryRollupOutcome {
25    pub drawer_id: String,
26    pub stats: IngestStats,
27}
28
29pub struct PreparedDiaryRollup {
30    pub drawer_id: String,
31    pub content: String,
32    pub room: String,
33    pub day: String,
34    pub stats: IngestStats,
35    pub importance: i32,
36    _lock_guard: lock::IngestLock,
37}
38
39pub fn current_rollup_day_utc() -> String {
40    let days_since_epoch = std::time::SystemTime::now()
41        .duration_since(std::time::UNIX_EPOCH)
42        .map(|duration| (duration.as_secs() / 86_400) as i64)
43        .unwrap_or(0);
44    let (year, month, day) = civil_from_days(days_since_epoch);
45    format!("{year:04}-{month:02}-{day:02}")
46}
47
48pub fn diary_rollup_drawer_id(room: &str, day: &str) -> String {
49    format!("drawer_agent-diary_{room}_day_{day}")
50}
51
52pub fn prepare_diary_rollup(
53    db: &Database,
54    content: &str,
55    wing: &str,
56    options: DiaryRollupOptions<'_>,
57) -> Result<PreparedDiaryRollup, IngestError> {
58    if wing != DIARY_ROLLUP_WING {
59        return Err(IngestError::DiaryRollupWrongWing {
60            wing: wing.to_string(),
61        });
62    }
63
64    let room = options
65        .room
66        .filter(|room| !room.trim().is_empty())
67        .ok_or(IngestError::DiaryRollupMissingRoom)?;
68    let day = options
69        .day
70        .map(ToOwned::to_owned)
71        .unwrap_or_else(current_rollup_day_utc);
72    let drawer_id = diary_rollup_drawer_id(room, &day);
73    let mut stats = IngestStats {
74        files: 1,
75        chunks: 1,
76        ..IngestStats::default()
77    };
78
79    let home = mempal_home_from_db(db);
80    let lock_key = format!("diary_rollup_{room}_{day}");
81    let lock_guard = lock::acquire_source_lock(&home, &lock_key, LOCK_TIMEOUT)?;
82    stats.lock_wait_ms = Some(lock_guard.wait_duration().as_millis() as u64);
83
84    let new_content =
85        match db
86            .get_drawer(&drawer_id)
87            .map_err(|source| IngestError::CheckDrawer {
88                drawer_id: drawer_id.clone(),
89                source,
90            })? {
91            Some(existing) => format!("{}{DIARY_ROLLUP_SEPARATOR}{}", existing.content, content),
92            None => content.to_string(),
93        };
94
95    if new_content.len() > DAILY_ROLLUP_LIMIT_BYTES {
96        return Err(IngestError::DailyRollupFull {
97            drawer_id,
98            limit_bytes: DAILY_ROLLUP_LIMIT_BYTES,
99            attempted_bytes: new_content.len(),
100        });
101    }
102
103    Ok(PreparedDiaryRollup {
104        drawer_id,
105        content: new_content,
106        room: room.to_string(),
107        day,
108        stats,
109        importance: options.importance,
110        _lock_guard: lock_guard,
111    })
112}
113
114pub fn commit_prepared_diary_rollup(
115    db: &Database,
116    prepared: PreparedDiaryRollup,
117    vector: &[f32],
118) -> Result<DiaryRollupOutcome, IngestError> {
119    let drawer = Drawer {
120        id: prepared.drawer_id.clone(),
121        content: prepared.content,
122        wing: DIARY_ROLLUP_WING.to_string(),
123        room: Some(prepared.room.clone()),
124        source_file: Some(format!(
125            "agent-diary://rollup/{}/{}",
126            prepared.room, prepared.day
127        )),
128        source_type: SourceType::Manual,
129        added_at: current_timestamp(),
130        chunk_index: Some(0),
131        normalize_version: CURRENT_NORMALIZE_VERSION,
132        importance: prepared.importance,
133        memory_kind: crate::core::types::MemoryKind::Evidence,
134        domain: MemoryDomain::Agent,
135        field: "diary".to_string(),
136        anchor_kind: AnchorKind::Repo,
137        anchor_id: "repo://agent-diary".to_string(),
138        parent_anchor_id: None,
139        provenance: Some(Provenance::Runtime),
140        statement: None,
141        tier: None,
142        status: None,
143        supporting_refs: Vec::new(),
144        counterexample_refs: Vec::new(),
145        teaching_refs: Vec::new(),
146        verification_refs: Vec::new(),
147        scope_constraints: None,
148        trigger_hints: None,
149    };
150
151    db.upsert_drawer_and_replace_vector(&drawer, vector)
152        .map_err(|source| IngestError::InsertDrawer {
153            drawer_id: prepared.drawer_id.clone(),
154            source,
155        })?;
156
157    Ok(DiaryRollupOutcome {
158        drawer_id: prepared.drawer_id,
159        stats: prepared.stats,
160    })
161}
162
163pub async fn ingest_diary_rollup<E: Embedder + ?Sized>(
164    db: &Database,
165    embedder: &E,
166    content: &str,
167    wing: &str,
168    options: DiaryRollupOptions<'_>,
169) -> Result<DiaryRollupOutcome, IngestError> {
170    if options.dry_run {
171        if wing != DIARY_ROLLUP_WING {
172            return Err(IngestError::DiaryRollupWrongWing {
173                wing: wing.to_string(),
174            });
175        }
176        let room = options
177            .room
178            .filter(|room| !room.trim().is_empty())
179            .ok_or(IngestError::DiaryRollupMissingRoom)?;
180        let day = options
181            .day
182            .map(ToOwned::to_owned)
183            .unwrap_or_else(current_rollup_day_utc);
184        let drawer_id = diary_rollup_drawer_id(room, &day);
185        let stats = IngestStats {
186            files: 1,
187            chunks: 1,
188            ..IngestStats::default()
189        };
190        return Ok(DiaryRollupOutcome { drawer_id, stats });
191    }
192
193    let prepared = prepare_diary_rollup(db, content, wing, options)?;
194
195    let vector = embedder
196        .embed(&[prepared.content.as_str()])
197        .await
198        .map_err(|source| IngestError::EmbedChunks {
199            path: std::path::PathBuf::from(format!(
200                "agent-diary://rollup/{}/{}",
201                prepared.room, prepared.day
202            )),
203            source,
204        })?
205        .into_iter()
206        .next()
207        .ok_or_else(|| IngestError::EmbedderReturnedNoVector {
208            drawer_id: prepared.drawer_id.clone(),
209        })?;
210
211    commit_prepared_diary_rollup(db, prepared, &vector)
212}
213
214fn civil_from_days(days_since_epoch: i64) -> (i32, u32, u32) {
215    let days = days_since_epoch + 719_468;
216    let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
217    let day_of_era = days - era * 146_097;
218    let year_of_era =
219        (day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
220    let mut year = year_of_era + era * 400;
221    let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
222    let month_prime = (5 * day_of_year + 2) / 153;
223    let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
224    let month = month_prime + if month_prime < 10 { 3 } else { -9 };
225    year += i64::from(month <= 2);
226    (year as i32, month as u32, day as u32)
227}