Follow up #205: survive dropped packs, validate pieces, budget beliefs

crates/tinymemory-integrations/src/cortex/README.md

@@ -178,7 +178,9 @@ as prefixes, so they cannot be labelled and are filtered only client-side.
   scopes, or, for an unscoped read, every kind scope the engine holds. With
   nothing to read the answer is empty and nothing is sent. The answer comes
   from the pack holding the most admitted events. The answer route is
-  called **once** with `use_pack_id`. Hosted omits a null
+  called **once** with `use_pack_id` (again, after every pack is built
+  anew, when the packs were dropped in between, up to three rounds:
+  CortexDB drops every pack on any forget). Hosted omits a null
   `answer_instructions`, because its schema is strict; Direct sends `null`.
   Citations come from the packs' decoded events, filtered (reach included),
   merged rank by rank (the most specific node's first), one per item, capped

crates/tinymemory-integrations/src/cortex/engine/beliefs.rs

@@ -34,7 +34,7 @@ use tinymemory_api::{
 };
 
 use super::CortexEngine;
-use super::fetch::recall_body;
+use super::fetch::{recall_body, whole_items_budget};
 use super::items::hit;
 use crate::cortex::descriptor::{CortexWire, Route};
 use crate::cortex::envelope::parse_scope;
@@ -174,6 +174,7 @@ impl CortexEngine {
         body["budgets"]["per_layer_limits"] = json!({
             "events": 0, "facts": 0, "episodes": 0, "understanding": 0, "beliefs": limit,
         });
+        body["budgets"]["max_tokens"] = json!(whole_items_budget(limit));
         let pack = self.log.recall(&body).await?;
         Ok(beliefs_in(&pack, "/layers/beliefs"))
     }

crates/tinymemory-integrations/src/cortex/engine/beliefs_tests.rs

@@ -98,6 +98,14 @@ async fn a_built_scope_s_beliefs_are_read_with_and_without_a_query() {
     let belief_pack = state.seen.lock().unwrap().recalls.last().cloned().unwrap();
     assert_eq!(belief_pack["include"], json!(["beliefs"]));
     assert_eq!(belief_pack["view"], "granular");
+    let asked = belief_pack["budgets"]["per_layer_limits"]["beliefs"]
+        .as_u64()
+        .unwrap();
+    assert_eq!(
+        belief_pack["budgets"]["max_tokens"].as_u64(),
+        Some(asked * crate::cortex::envelope::chunks::MAX_EVENT_TEXT_BYTES as u64),
+        "room for each belief asked for"
+    );
     assert_eq!(ranked[0].meta.namespace, node);
     let listed = engine
         .beliefs(BeliefsRequest::new(Reach::exact(node), 5))

crates/tinymemory-integrations/src/cortex/engine/mod_tests.rs

@@ -320,3 +320,46 @@ async fn a_store_succeeds_when_ranked_recall_is_down() {
         assert_eq!(listed.items.len(), 1, "the settle probe is best-effort");
     }
 }
+
+#[tokio::test]
+async fn an_answer_whose_pack_expired_recalls_that_scope_again() {
+    for (engine, state) in both().await {
+        for item in items() {
+            engine.store(item).await.unwrap();
+        }
+        let mut req = RecallRequest::new("which editor helix", 2);
+        req.filter = MetaFilter::kinds([ItemKind::Learning]);
+        state.seen.lock().unwrap().recalls.clear();
+        state.expire_packs.store(1, Ordering::SeqCst);
+        let answer = engine.recall(req.clone()).await.unwrap();
+        assert_eq!(answer.answer, "grounded answer for which editor helix");
+        {
+            let seen = state.seen.lock().unwrap();
+            assert_eq!(seen.recalls.len(), 2, "the one scope, packed again");
+            assert_eq!(seen.recalls[0], seen.recalls[1], "the same pack request");
+            assert_eq!(seen.answers.len(), 2, "answered from the fresh pack");
+        }
+
+        // With several scopes every pack is read again, so the citations
+        // come from packs read after whatever dropped them, not before.
+        let mut wide = RecallRequest::new("which editor helix", 2);
+        wide.filter = MetaFilter::kinds([ItemKind::Learning, ItemKind::Conversation]);
+        state.seen.lock().unwrap().recalls.clear();
+        state.expire_packs.store(1, Ordering::SeqCst);
+        engine.recall(wide).await.unwrap();
+        assert_eq!(
+            state.seen.lock().unwrap().recalls.len(),
+            4,
+            "two scopes, each packed in both rounds"
+        );
+
+        state.seen.lock().unwrap().answers.clear();
+        state.expire_packs.store(3, Ordering::SeqCst);
+        assert!(
+            matches!(engine.recall(req).await, Err(Error::NotFound(_))),
+            "a bounded retry, not forever"
+        );
+        assert_eq!(state.seen.lock().unwrap().answers.len(), 3);
+        state.expire_packs.store(0, Ordering::SeqCst);
+    }
+}

crates/tinymemory-integrations/src/cortex/engine/recall.rs

@@ -9,7 +9,9 @@
 //! order, not by relevance, so it would answer from an arbitrary sample.
 //! With no scope to read, the answer is empty and nothing is sent.
 //!
-//! The answer route is asked once with `use_pack_id`, so it answers from
+//! The answer route is asked once with `use_pack_id` (again, after every
+//! pack is built anew, when CortexDB dropped the packs in between, up to
+//! three rounds: it drops every pack on any forget), so it answers from
 //! exactly the evidence that pack holds: with several packs, the one holding
 //! the most admitted events, the most specific node on a tie. **The answer
 //! text is grounded on that one pack, while the citations come from every
@@ -43,6 +45,10 @@ use crate::cortex::error::{Error, Result};
 /// Recall packs built at once when a reach spans several scopes.
 const PACKS_AT_ONCE: usize = 4;
 
+/// Rounds (packs, then one answer) at most per recall: the first, and one
+/// after each time CortexDB dropped the packs before the answer used one.
+const ANSWER_ATTEMPTS: usize = 3;
+
 /// The derived layers a pack also draws on, besides events.
 const DERIVED_LAYERS: [&str; 4] = ["facts", "beliefs", "episodes", "understanding"];
 
@@ -71,6 +77,13 @@ fn pack_budgets(limit: usize) -> Value {
     Value::Object(layers)
 }
 
+/// The id a recall pack names.
+fn pack_id_of(pack: &Value) -> Result<&str> {
+    pack.get("pack_id")
+        .and_then(Value::as_str)
+        .ok_or_else(|| Error::Engine("CortexDB recall omitted pack_id".to_string()))
+}
+
 /// The answer request body.
 ///
 /// The hosted route's schema is strict (an unknown key, or a `null`
@@ -120,39 +133,20 @@ impl CortexEngine {
         let mut ordered: Vec<&KindScope> = scopes.iter().collect();
         ordered.sort_by_key(|scope| std::cmp::Reverse(scope.namespace.depth()));
         let paths: Vec<String> = ordered.into_iter().map(|s| s.path.clone()).collect();
-        let req = &req;
-        let packs: Vec<(String, Value)> = stream::iter(paths)
-            .map(|path| async move {
-                let pack = self.pack(req, &path).await?;
-                Ok::<_, Error>((path, pack))
-            })
-            .buffered(PACKS_AT_ONCE)
-            .try_collect()
-            .await?;
-        let per_pack: Vec<Vec<Envelope>> = packs
-            .iter()
-            .map(|(_, pack)| ranked(pack, None, &req.filter))
-            .collect();
-        let chosen = per_pack
-            .iter()
-            .enumerate()
-            .max_by_key(|(index, events)| (events.len(), std::cmp::Reverse(*index)))
-            .map_or(0, |(index, _)| index);
-        let (scope, pack) = &packs[chosen];
-        let pack_id = pack
-            .get("pack_id")
-            .and_then(Value::as_str)
-            .ok_or_else(|| Error::Engine("CortexDB recall omitted pack_id".to_string()))?;
-        let answered = self
-            .log
-            .answer(&answer_body(
-                self.wire(),
-                scope,
-                &req.question,
-                pack_id,
-                req.instructions.as_deref(),
-            ))
-            .await?;
+        // A pack lives 60 s, and CortexDB drops every pack it holds once
+        // anything is forgotten (measured on 0.10.4: a forget in another
+        // scope turns the next `use_pack_id` into a 404). On a 404 every
+        // pack is built again, so the answer and the citations both come
+        // from packs read after that forget, up to [`ANSWER_ATTEMPTS`]
+        // rounds in all.
+        let mut round = 1;
+        let (per_pack, answered) = loop {
+            let (per_pack, answered) = self.answer_round(&req, &paths).await?;
+            match answered {
+                Err(Error::NotFound(_)) if round < ANSWER_ATTEMPTS => round += 1,
+                answered => break (per_pack, answered?),
+            }
+        };
         let answer = answered
             .get("answer")
             .and_then(Value::as_str)
@@ -183,6 +177,43 @@ impl CortexEngine {
 }
 
 impl CortexEngine {
+    /// One round of recall: a pack per path (four at a time), decoded and
+    /// filtered, and the answer route asked once with the pack holding the
+    /// most admitted events, the most specific node on a tie. The answer's
+    /// own error is handed back, so the caller can tell a dropped pack.
+    async fn answer_round(
+        &self,
+        req: &RecallRequest,
+        paths: &[String],
+    ) -> Result<(Vec<Vec<Envelope>>, Result<Value>)> {
+        let packs: Vec<(String, Value)> = stream::iter(paths.to_vec())
+            .map(|path| async move {
+                let pack = self.pack(req, &path).await?;
+                Ok::<_, Error>((path, pack))
+            })
+            .buffered(PACKS_AT_ONCE)
+            .try_collect()
+            .await?;
+        let per_pack: Vec<Vec<Envelope>> = packs
+            .iter()
+            .map(|(_, pack)| ranked(pack, None, &req.filter))
+            .collect();
+        let chosen = per_pack
+            .iter()
+            .enumerate()
+            .max_by_key(|(index, events)| (events.len(), std::cmp::Reverse(*index)))
+            .map_or(0, |(index, _)| index);
+        let (scope, pack) = &packs[chosen];
+        let body = answer_body(
+            self.wire(),
+            scope,
+            &req.question,
+            pack_id_of(pack)?,
+            req.instructions.as_deref(),
+        );
+        Ok((per_pack, self.log.answer(&body).await))
+    }
+
     /// A recall pack for `req` over exactly `scope`, sized for `req.limit`
     /// citations. It includes the derived layers the answer route reads,
     /// after the events the citations come from.

crates/tinymemory-integrations/src/cortex/envelope/mod_tests.rs

@@ -287,3 +287,51 @@ fn a_whitespace_only_document_keeps_its_text() {
         assert_eq!(rebuild(&envelopes).unwrap(), item, "{body:?}");
     }
 }
+
+#[test]
+fn a_whole_read_needs_every_piece_of_one_agreed_layout() {
+    let item = long_document();
+    let pieces = Envelope::for_item(&item, "id").unwrap();
+    assert!(pieces.len() > 1);
+    assert_eq!(rebuild_whole(&pieces), Some(item.clone()));
+    assert_eq!(rebuild_whole(&pieces[1..]), None, "a piece missing");
+
+    let relaid = |edit: &dyn Fn(&mut ChunkInfo)| {
+        let mut changed = pieces.clone();
+        edit(changed[0].chunk.as_mut().unwrap());
+        rebuild_whole(&changed)
+    };
+    assert_eq!(relaid(&|chunk| chunk.count += 1), None, "counts disagree");
+    assert_eq!(
+        relaid(&|chunk| chunk.index = 99),
+        None,
+        "an index past the count"
+    );
+    let mut zero = pieces[..1].to_vec();
+    zero[0].chunk = Some(ChunkInfo {
+        index: 0,
+        count: 0,
+        pages: None,
+        section: None,
+    });
+    assert_eq!(rebuild_whole(&zero), None, "a zero count");
+
+    let mut whole = Envelope::for_item(&StoreItem::document("x", meta()), "id")
+        .unwrap()
+        .remove(0);
+    whole.text = match &item {
+        StoreItem::Document {
+            body: DocumentBody::Text(text),
+            ..
+        } => text.clone(),
+        _ => unreachable!("a document"),
+    };
+    whole.title.clone_from(&pieces[0].title);
+    whole.meta = pieces[0].meta.clone();
+    let mixed = vec![pieces[1].clone(), whole];
+    assert_eq!(
+        rebuild_whole(&mixed),
+        Some(item),
+        "the same item written whole before chunking reads as that body"
+    );
+}

crates/tinymemory-integrations/src/cortex/envelope/rebuild.rs

@@ -83,35 +83,52 @@ pub(crate) fn rebuild(envelopes: &[Envelope]) -> Option<StoreItem> {
 /// does not list yet. A `fetch` hit or a `recall` citation is one piece and
 /// uses [`rebuild`].
 pub(crate) fn rebuild_whole(envelopes: &[Envelope]) -> Option<StoreItem> {
-    if let Some(count) = envelopes.iter().find_map(|e| Some(e.chunk.as_ref()?.count)) {
-        let held: HashSet<u32> = envelopes
-            .iter()
-            .filter_map(|envelope| Some(envelope.chunk.as_ref()?.index))
-            .collect();
-        if (0..count).any(|index| !held.contains(&index)) {
-            log::debug!(
-                "[cortex] chunked document {:?} is missing pieces; not returned whole",
-                envelopes[0].id
-            );
-            return None;
-        }
+    if !pieces_complete(envelopes) {
+        log::debug!(
+            "[cortex] chunked document {:?} is missing pieces or disagrees on its layout; \
+             not returned whole",
+            envelopes.first().map(|envelope| &envelope.id)
+        );
+        return None;
     }
     rebuild(envelopes)
 }
 
-/// A document's text from its envelopes: the first one's, or the pieces of
-/// a chunked document in index order, each once.
+/// Whether `envelopes` hold a whole item: one written as a single event
+/// (a document's unchunked envelope is its whole body), or every piece of a
+/// chunked document, all agreeing on one positive count with every index
+/// below it and each index present.
+fn pieces_complete(envelopes: &[Envelope]) -> bool {
+    if envelopes.iter().any(|envelope| envelope.chunk.is_none()) {
+        return true;
+    }
+    let Some(count) = envelopes.iter().find_map(|e| Some(e.chunk.as_ref()?.count)) else {
+        return true;
+    };
+    let mut held = HashSet::new();
+    for chunk in envelopes
+        .iter()
+        .filter_map(|envelope| envelope.chunk.as_ref())
+    {
+        if chunk.count != count || chunk.index >= count {
+            return false;
+        }
+        held.insert(chunk.index);
+    }
+    count > 0 && held.len() == count as usize
+}
+
+/// A document's text from its envelopes: an unchunked envelope's (the whole
+/// body: the same item written before chunking, since an id is a content
+/// digest), else the pieces of a chunked document in index order, each once.
 fn document_text(envelopes: &[Envelope]) -> String {
+    if let Some(whole) = envelopes.iter().find(|envelope| envelope.chunk.is_none()) {
+        return whole.text.clone();
+    }
     let mut pieces: Vec<(u32, &str)> = envelopes
         .iter()
         .filter_map(|envelope| Some((envelope.chunk.as_ref()?.index, envelope.text.as_str())))
         .collect();
-    if pieces.is_empty() {
-        return envelopes
-            .first()
-            .map(|envelope| envelope.text.clone())
-            .unwrap_or_default();
-    }
     pieces.sort_by_key(|(index, _)| *index);
     pieces.dedup_by_key(|(index, _)| *index);
     pieces.into_iter().map(|(_, text)| text).collect()

crates/tinymemory-integrations/src/cortex/testing/mod.rs

@@ -77,6 +77,9 @@ pub(crate) struct Double {
     pub(crate) rate_limit_forget: AtomicUsize,
     /// Recall answers 500.
     pub(crate) recall_down: AtomicBool,
+    /// Answers refuse their `use_pack_id` as expired (404) this many times,
+    /// as CortexDB does once anything is forgotten after the pack was built.
+    pub(crate) expire_packs: AtomicUsize,
     /// Once the next write is applied, rate limit this many listings and
     /// hide the listing this many more times: (429s, hidden). Lets a test
     /// aim at the reads a write makes after it is sent, not the replay

crates/tinymemory-integrations/src/cortex/testing/routes.rs

@@ -297,6 +297,9 @@ async fn answer(
     if body["use_pack_id"].as_str() != Some("pack_test") {
         return fail(&state, 400, "MISSING_PACK");
     }
+    if take_one(&state.expire_packs) {
+        return fail(&state, 404, "NOT_FOUND");
+    }
     ok(
         &state,
         200,