writerslogic_warm/core/documents/
listing.rs1use serde::Serialize;
7use serde_json::Value;
8
9use super::index::text_of;
10use super::text::TextSegment;
11use super::{DocumentStore, StoredDocument};
12
13#[derive(Clone, Debug, Serialize)]
15#[serde(rename_all = "camelCase")]
16pub struct DocumentRecord {
17 pub id: String,
18 pub source_uri: Option<String>,
19 pub version: String,
20 pub metadata: Value,
21 pub text: Option<String>,
24 pub chunk_count: usize,
25}
26
27fn record(document: &StoredDocument, segment: Option<&TextSegment>) -> DocumentRecord {
28 DocumentRecord {
29 id: document.id.clone(),
30 source_uri: document.source_uri.clone(),
31 version: document.version.clone(),
32 metadata: document.metadata.clone(),
33 text: document
34 .chunks
35 .first()
36 .and_then(|chunk| text_of(chunk, segment)),
37 chunk_count: document.chunks.len(),
38 }
39}
40
41pub(crate) fn get_record(store: &DocumentStore, id: &str) -> Option<DocumentRecord> {
42 store.get(id).map(|document| {
43 let segment = store.text_segment();
44 record(document, segment.as_deref())
45 })
46}
47
48pub(crate) fn list_records(
50 store: &DocumentStore,
51 from: Option<&str>,
52 limit: usize,
53) -> Vec<DocumentRecord> {
54 use std::ops::Bound::{Included, Unbounded};
55 let start = from.map_or(Unbounded, Included);
56 let segment = store.text_segment();
57 store
58 .documents
59 .range::<str, _>((start, Unbounded))
60 .take(limit)
61 .map(|(_, document)| record(document, segment.as_deref()))
62 .collect()
63}
64
65pub(crate) fn count(store: &DocumentStore) -> usize {
66 store.documents.len()
67}
68
69#[cfg(test)]
70mod tests {
71 use super::super::{prepare, DocumentInput, DocumentStore};
72 use super::*;
73
74 fn store_with(ids: &[&str]) -> DocumentStore {
75 let mut store = DocumentStore::default();
76 for id in ids {
77 let input = DocumentInput {
78 id: (*id).into(),
79 text: format!("body of {id}"),
80 source_uri: None,
81 version: Some("v1".into()),
82 metadata: Some(serde_json::json!({"k": 1})),
83 chunk_words: None,
84 overlap_words: None,
85 store_text: Some(true),
86 embeddings: None,
87 title_weight: None,
88 };
89 let prepared = prepare(input, 256, None, Default::default()).unwrap();
90 store.insert(prepared.id.clone(), prepared);
91 }
92 store
93 }
94
95 #[test]
96 fn lists_in_id_order_and_resumes_after_a_cursor() {
97 let store = store_with(&["b", "a", "d", "c"]);
98 assert_eq!(count(&store), 4);
99 let first = list_records(&store, None, 2);
100 let ids =
101 |records: &[DocumentRecord]| records.iter().map(|r| r.id.clone()).collect::<Vec<_>>();
102 assert_eq!(ids(&first), ["a", "b"]);
103 let resume = format!("{}\0", first[1].id);
104 assert_eq!(ids(&list_records(&store, Some(&resume), 10)), ["c", "d"]);
105 assert!(list_records(&store, Some("e"), 10).is_empty());
106 assert_eq!(list_records(&store, Some("c"), 1)[0].id, "c");
107 }
108
109 #[test]
110 fn get_returns_text_metadata_and_version() {
111 let store = store_with(&["x"]);
112 let record = get_record(&store, "x").unwrap();
113 assert_eq!(record.text.as_deref(), Some("body of x"));
114 assert_eq!(record.version, "v1");
115 assert_eq!(record.metadata["k"], 1);
116 assert_eq!(record.chunk_count, 1);
117 assert!(get_record(&store, "missing").is_none());
118 }
119}