Skip to main content

writerslogic_warm/core/documents/
listing.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3//! Read access to stored documents by id and in id order, for callers that must enumerate or
4//! export a store (the search path only returns documents that match a query).
5
6use serde::Serialize;
7use serde_json::Value;
8
9use super::index::text_of;
10use super::text::TextSegment;
11use super::{DocumentStore, StoredDocument};
12
13/// A stored document as the store holds it: first chunk's text, metadata and version.
14#[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    /// Text of the first chunk; `None` for a document stored without text. A document that
22    /// fit one chunk when stored reads back exactly (modulo leading and trailing whitespace).
23    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
48/// Up to `limit` documents with ids at or after `from`, in id order.
49pub(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}