gemini_memory_rs/okf/
document.rs

1//! The OKF document container: YAML front matter plus Markdown body sections.
2//!
3//! A memory record is a Markdown file a human can read, diff and hand-edit.
4//! The front matter carries machine state; the body carries the sentences the
5//! model will actually see.
6
7use std::collections::BTreeMap;
8
9use super::yaml::{self, Yaml};
10use crate::core::MemoryError;
11
12/// The front-matter delimiter.
13const FENCE: &str = "---";
14
15/// A parsed OKF file.
16#[derive(Debug, Clone, PartialEq)]
17pub struct OkfDocument {
18    /// Machine-readable front matter.
19    pub front: Yaml,
20    /// Body sections, keyed by their `#` heading text.
21    pub sections: BTreeMap<String, String>,
22    /// Heading order as it appeared, so emitting is stable.
23    pub section_order: Vec<String>,
24}
25
26impl OkfDocument {
27    /// Build a document from front matter and ordered sections.
28    pub fn new(front: Yaml, sections: Vec<(&str, String)>) -> Self {
29        let mut order = Vec::new();
30        let mut map = BTreeMap::new();
31        for (heading, body) in sections {
32            order.push(heading.to_string());
33            map.insert(heading.to_string(), body);
34        }
35        Self {
36            front,
37            sections: map,
38            section_order: order,
39        }
40    }
41
42    /// Read a section body by heading.
43    pub fn section(&self, heading: &str) -> Option<&str> {
44        self.sections.get(heading).map(|s| s.trim())
45    }
46
47    /// Read a section as a `- item` list.
48    pub fn section_list(&self, heading: &str) -> Vec<String> {
49        self.section(heading)
50            .map(|body| {
51                body.lines()
52                    .filter_map(|line| line.trim().strip_prefix("- ").map(str::to_string))
53                    .collect()
54            })
55            .unwrap_or_default()
56    }
57
58    /// Parse an OKF file.
59    pub fn parse(source: &str, path: &str) -> Result<Self, MemoryError> {
60        let source = source.trim_start_matches('\u{feff}');
61        let rest = source
62            .strip_prefix(FENCE)
63            .and_then(|r| r.strip_prefix('\n').or_else(|| r.strip_prefix("\r\n")))
64            .ok_or_else(|| MemoryError::MalformedRecord {
65                path: path.to_string(),
66                message: "file does not begin with a `---` front-matter fence".to_string(),
67            })?;
68
69        let end = rest
70            .find("\n---")
71            .ok_or_else(|| MemoryError::MalformedRecord {
72                path: path.to_string(),
73                message: "front matter is not terminated by `---`".to_string(),
74            })?;
75
76        let front_src = &rest[..end];
77        let body = rest[end + 4..].trim_start_matches(['\r', '\n']);
78
79        let front = yaml::parse(front_src).map_err(|e| MemoryError::MalformedRecord {
80            path: path.to_string(),
81            message: e.to_string(),
82        })?;
83
84        let (sections, section_order) = parse_sections(body);
85        Ok(Self {
86            front,
87            sections,
88            section_order,
89        })
90    }
91
92    /// Parse a category file holding several records back to back.
93    ///
94    /// Records are delimited by `---` fences in pairs: open, front matter,
95    /// close, body, then the next record's open fence. A body line that is
96    /// exactly `---` would therefore be read as a record boundary, so
97    /// [`Self::render_many`] refuses to write one.
98    pub fn parse_many(source: &str, path: &str) -> Result<Vec<Self>, MemoryError> {
99        let source = source.trim_start_matches('\u{feff}');
100        let lines: Vec<&str> = source.lines().collect();
101        let fences: Vec<usize> = lines
102            .iter()
103            .enumerate()
104            .filter(|(_, l)| l.trim_end() == FENCE)
105            .map(|(i, _)| i)
106            .collect();
107
108        if fences.is_empty() {
109            return if source.trim().is_empty() {
110                Ok(Vec::new())
111            } else {
112                Err(MemoryError::MalformedRecord {
113                    path: path.to_string(),
114                    message: "file contains no `---` front-matter fence".to_string(),
115                })
116            };
117        }
118        if !fences.len().is_multiple_of(2) {
119            return Err(MemoryError::MalformedRecord {
120                path: path.to_string(),
121                message: format!(
122                    "odd number of `---` fences ({}) — a record is unterminated",
123                    fences.len()
124                ),
125            });
126        }
127
128        let mut docs = Vec::new();
129        for pair in 0..fences.len() / 2 {
130            let open = fences[pair * 2];
131            let close = fences[pair * 2 + 1];
132            let body_end = fences.get((pair + 1) * 2).copied().unwrap_or(lines.len());
133            let mut chunk = String::new();
134            for line in &lines[open..close.min(body_end)] {
135                chunk.push_str(line);
136                chunk.push('\n');
137            }
138            chunk.push_str(FENCE);
139            chunk.push('\n');
140            for line in &lines[close + 1..body_end] {
141                chunk.push_str(line);
142                chunk.push('\n');
143            }
144            docs.push(Self::parse(&chunk, path)?);
145        }
146        Ok(docs)
147    }
148
149    /// Render several records into one category file.
150    pub fn render_many(docs: &[Self], path: &str) -> Result<String, MemoryError> {
151        let mut out = String::new();
152        for (idx, doc) in docs.iter().enumerate() {
153            if doc
154                .sections
155                .values()
156                .any(|body| body.lines().any(|l| l.trim_end() == FENCE))
157            {
158                return Err(MemoryError::MalformedRecord {
159                    path: path.to_string(),
160                    message: "record body contains a `---` line, which would be read as a \
161                              record boundary"
162                        .to_string(),
163                });
164            }
165            if idx > 0 {
166                out.push('\n');
167            }
168            out.push_str(&doc.to_markdown());
169        }
170        Ok(out)
171    }
172
173    /// Render the document back to text.
174    pub fn to_markdown(&self) -> String {
175        let mut out = String::new();
176        out.push_str(FENCE);
177        out.push('\n');
178        out.push_str(&yaml::emit(&self.front));
179        out.push_str(FENCE);
180        out.push('\n');
181        for heading in &self.section_order {
182            if let Some(body) = self.sections.get(heading) {
183                out.push_str("\n# ");
184                out.push_str(heading);
185                out.push('\n');
186                let trimmed = body.trim();
187                if !trimmed.is_empty() {
188                    out.push_str(trimmed);
189                    out.push('\n');
190                }
191            }
192        }
193        out
194    }
195}
196
197fn parse_sections(body: &str) -> (BTreeMap<String, String>, Vec<String>) {
198    let mut sections = BTreeMap::new();
199    let mut order = Vec::new();
200    let mut current: Option<String> = None;
201    let mut buffer = String::new();
202
203    for line in body.lines() {
204        if let Some(heading) = line.strip_prefix("# ") {
205            if let Some(name) = current.take() {
206                sections.insert(name, buffer.trim().to_string());
207            }
208            buffer = String::new();
209            let heading = heading.trim().to_string();
210            order.push(heading.clone());
211            current = Some(heading);
212        } else if current.is_some() {
213            buffer.push_str(line);
214            buffer.push('\n');
215        }
216    }
217    if let Some(name) = current {
218        sections.insert(name, buffer.trim().to_string());
219    }
220    (sections, order)
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    const SAMPLE: &str = r#"---
228okf: memory/v1
229id: mem_01K4D8P3
230confidence: 1.0
231retrieval:
232  tags:
233    - food
234    - diet
235---
236# Fact
237The user is pescatarian.
238
239# Evidence Summary
240Explicitly stated by the user.
241
242# Supersedes
243- mem_01JVEGETARIAN
244"#;
245
246    #[test]
247    fn parses_front_matter_and_sections() {
248        let doc = OkfDocument::parse(SAMPLE, "sample.md").unwrap();
249        assert_eq!(doc.front.get("id").unwrap().as_str(), Some("mem_01K4D8P3"));
250        assert_eq!(doc.section("Fact"), Some("The user is pescatarian."));
251        assert_eq!(doc.section_list("Supersedes"), vec!["mem_01JVEGETARIAN"]);
252    }
253
254    #[test]
255    fn round_trips_to_markdown() {
256        let doc = OkfDocument::parse(SAMPLE, "sample.md").unwrap();
257        let rendered = doc.to_markdown();
258        let reparsed = OkfDocument::parse(&rendered, "rendered.md").unwrap();
259        assert_eq!(doc.front, reparsed.front);
260        assert_eq!(doc.sections, reparsed.sections);
261    }
262
263    #[test]
264    fn a_missing_fence_names_the_file() {
265        let err = OkfDocument::parse("no front matter", "records/x.md").unwrap_err();
266        match err {
267            MemoryError::MalformedRecord { path, .. } => assert_eq!(path, "records/x.md"),
268            other => panic!("unexpected error: {other}"),
269        }
270    }
271
272    #[test]
273    fn an_unterminated_fence_is_rejected() {
274        assert!(OkfDocument::parse("---\nid: x\n", "x.md").is_err());
275    }
276
277    #[test]
278    fn category_files_hold_several_records() {
279        let first = OkfDocument::parse(SAMPLE, "s.md").unwrap();
280        let second = OkfDocument::new(
281            yaml::parse("okf: memory/v1\nid: mem_two\n").unwrap(),
282            vec![("Fact", "The user drinks flat whites.".to_string())],
283        );
284        let file = OkfDocument::render_many(&[first.clone(), second.clone()], "cat.md").unwrap();
285
286        let parsed = OkfDocument::parse_many(&file, "cat.md").unwrap();
287        assert_eq!(parsed.len(), 2);
288        assert_eq!(parsed[0].front, first.front);
289        assert_eq!(parsed[1].section("Fact"), second.section("Fact"));
290    }
291
292    #[test]
293    fn an_empty_category_file_parses_to_no_records() {
294        assert!(OkfDocument::parse_many("", "cat.md").unwrap().is_empty());
295        assert!(
296            OkfDocument::parse_many("\n\n", "cat.md")
297                .unwrap()
298                .is_empty()
299        );
300    }
301
302    #[test]
303    fn a_body_containing_a_fence_is_refused_at_write_time() {
304        let doc = OkfDocument::new(
305            yaml::parse("okf: memory/v1\n").unwrap(),
306            vec![("Fact", "before\n---\nafter".to_string())],
307        );
308        assert!(OkfDocument::render_many(&[doc], "cat.md").is_err());
309    }
310
311    #[test]
312    fn an_odd_fence_count_names_the_problem() {
313        let err =
314            OkfDocument::parse_many("---\nid: a\n---\n# Fact\nx\n---\n", "cat.md").unwrap_err();
315        assert!(err.to_string().contains("unterminated"));
316    }
317}
gemini_memory_rs/okf/document.rs

gemini_memory_rs/okf/
document.rs
