Skip to main content

gonzalo_knowledge/
lib.rs

1//! Knowledge-store capability for gonzalo.
2//!
3//! A single "what do we know about X" surface that composes the existing
4//! capability layers (ADR 0011): a [`Store`] for records, a [`VectorIndex`] for
5//! semantic retrieval, and an [`Embedder`] for turning text into vectors — all
6//! addressed by the shared [`RecordKey`](gonzalo_core::RecordKey). Queries return first-class
7//! [`Record`]s as [`Hit`]s, not bare ids.
8//!
9//! Which record kinds are knowledge-bearing, and how their text is split into
10//! embeddable pieces, lives in [`chunk`] (with [`knowledge_text`] as the joined
11//! one-document view). Records are embedded at chunk granularity — a `Session`
12//! by turn, a long `MemoryTier`/`Ticket` by paragraph, a `Topic` by bullet — so
13//! a query matching one turn/section isn't drowned out by a document average;
14//! [`KnowledgeStore::query`] de-dups chunk matches back to one hit per record.
15//! A `graph`-backed join ([`KnowledgeStore::query_in_subgraph`]) is a further
16//! refinement behind this same surface.
17
18use gonzalo_core::{KeyPrefix, Record, RecordKind, Result, Store};
19use gonzalo_domain::{MemoryTier, RecordCodec, Session, Ticket, TicketEvent, Topic};
20use gonzalo_vector::{Embedder, VectorIndex};
21
22/// One search hit: a first-class record and its similarity score (higher is
23/// more similar).
24#[derive(Debug, Clone)]
25pub struct Hit {
26    pub record: Record,
27    pub score: f32,
28}
29
30/// One chunk-granular search hit: the parent [`Record`], the matching chunk's
31/// score, its `ordinal` within the record, and that chunk's `text`. Unlike a
32/// [`Hit`], chunk hits are **not** de-duped to one-per-record.
33#[derive(Debug, Clone)]
34pub struct ChunkHit {
35    pub record: Record,
36    pub score: f32,
37    pub ordinal: usize,
38    pub text: String,
39}
40
41/// Composes a [`Store`], a [`VectorIndex`], and an [`Embedder`] into one
42/// retrieval surface keyed by [`RecordKey`](gonzalo_core::RecordKey).
43pub struct KnowledgeStore<S, V, E> {
44    store: S,
45    index: V,
46    embedder: E,
47    /// Last-seen chunk count per record, so a re-ingest that shrinks a record
48    /// can remove the now-orphaned high-ordinal chunks from the index. Held
49    /// in-memory, matching the (currently in-memory) index's lifecycle — see the
50    /// design note in the module docs.
51    chunk_counts: std::sync::Mutex<std::collections::HashMap<gonzalo_core::RecordKey, usize>>,
52}
53
54impl<S: Store, V: VectorIndex, E: Embedder> KnowledgeStore<S, V, E> {
55    pub fn new(store: S, index: V, embedder: E) -> Self {
56        Self {
57            store,
58            index,
59            embedder,
60            chunk_counts: std::sync::Mutex::new(std::collections::HashMap::new()),
61        }
62    }
63
64    /// Borrow the underlying store (e.g. to put records before ingesting them).
65    pub fn store(&self) -> &S {
66        &self.store
67    }
68
69    /// Ingest the record at `key`: split it into per-kind [`chunk`]s, embed each,
70    /// and index it under a derived per-chunk key. Returns `Ok(false)` (without
71    /// indexing) if the record is absent or its kind is definitionally not
72    /// knowledge-bearing; returns `Err` if a knowledge-bearing body fails to
73    /// parse — a corrupt record is a real failure, not "not indexable" (#139).
74    pub async fn ingest(&self, key: &gonzalo_core::RecordKey) -> Result<bool> {
75        let Some(record) = self.store.get(key).await? else {
76            return Ok(false);
77        };
78        let Some(chunks) = chunk(&record)? else {
79            return Ok(false);
80        };
81
82        // Remove chunks orphaned by a shrink since the last ingest. Read the old
83        // count under a scoped lock — never held across an `.await`.
84        let old = {
85            let counts = self.chunk_counts.lock().unwrap();
86            counts.get(key).copied().unwrap_or(0)
87        };
88        for ordinal in chunks.len()..old {
89            self.index.remove(&chunk_key(key, ordinal)).await?;
90        }
91
92        for (ordinal, text) in chunks.iter().enumerate() {
93            let vector = self.embedder.embed(text).await?;
94            self.index.upsert(chunk_key(key, ordinal), vector).await?;
95        }
96
97        self.chunk_counts
98            .lock()
99            .unwrap()
100            .insert(key.clone(), chunks.len());
101        Ok(true)
102    }
103
104    /// De-index the record at `key`: remove every chunk vector it contributed to
105    /// the index and forget its chunk count. This is the counterpart to a record
106    /// being deleted from the store — without it, the record's chunk vectors are
107    /// orphaned, still matching in [`query`] but de-duping to a parent that no
108    /// longer resolves, so they silently shrink the result set (#150).
109    ///
110    /// Idempotent: removing a record that was never ingested (chunk count 0)
111    /// removes nothing. Uses the same [`chunk_key`] derivation as [`ingest`].
112    pub async fn remove(&self, key: &gonzalo_core::RecordKey) -> Result<()> {
113        // Read the count under a scoped lock — never held across an `.await`.
114        let count = {
115            let counts = self.chunk_counts.lock().unwrap();
116            counts.get(key).copied().unwrap_or(0)
117        };
118        for ordinal in 0..count {
119            self.index.remove(&chunk_key(key, ordinal)).await?;
120        }
121        self.chunk_counts.lock().unwrap().remove(key);
122        Ok(())
123    }
124
125    /// Semantic query, one hit per record. Embed `text`, over-fetch chunk
126    /// matches (restricted to `filter`), collapse them to their parent records
127    /// keeping each record's best chunk score, and return the top-`k` records.
128    ///
129    /// Because a record with many matching chunks can crowd the over-fetch
130    /// window, pathological cases may return fewer than `k` hits.
131    pub async fn query(&self, text: &str, k: usize, filter: &KeyPrefix) -> Result<Vec<Hit>> {
132        if k == 0 {
133            return Ok(Vec::new());
134        }
135        let query_vec = self.embedder.embed(text).await?;
136        let fetch = k.saturating_mul(OVERFETCH);
137        let matches = self.index.query(&query_vec, fetch, filter).await?;
138
139        // Collapse chunk matches to parents, keeping the best score per parent.
140        let mut best: std::collections::BTreeMap<gonzalo_core::RecordKey, f32> =
141            std::collections::BTreeMap::new();
142        for m in matches {
143            let (parent, _ordinal) = parent_key(&m.key);
144            best.entry(parent)
145                .and_modify(|s| *s = s.max(m.score))
146                .or_insert(m.score);
147        }
148
149        // Rank parents by best chunk score, resolve the top-k to records.
150        let mut ranked: Vec<(gonzalo_core::RecordKey, f32)> = best.into_iter().collect();
151        ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
152        ranked.truncate(k);
153
154        let mut hits = Vec::with_capacity(ranked.len());
155        for (parent, score) in ranked {
156            if let Some(record) = self.store.get(&parent).await? {
157                hits.push(Hit { record, score });
158            }
159        }
160        Ok(hits)
161    }
162
163    /// Chunk-granular query: the top-`k` matching chunks (restricted to
164    /// `filter`), **not** de-duped to their parents. Each hit carries the parent
165    /// record, the chunk's ordinal, and the chunk's text.
166    pub async fn query_chunks(
167        &self,
168        text: &str,
169        k: usize,
170        filter: &KeyPrefix,
171    ) -> Result<Vec<ChunkHit>> {
172        if k == 0 {
173            return Ok(Vec::new());
174        }
175        let query_vec = self.embedder.embed(text).await?;
176        let matches = self.index.query(&query_vec, k, filter).await?;
177        let mut hits = Vec::with_capacity(matches.len());
178        for m in matches {
179            let (parent, ordinal) = parent_key(&m.key);
180            if let Some(record) = self.store.get(&parent).await? {
181                let text = chunk(&record)?
182                    .and_then(|cs| cs.into_iter().nth(ordinal))
183                    .unwrap_or_default();
184                hits.push(ChunkHit {
185                    record,
186                    score: m.score,
187                    ordinal,
188                    text,
189                });
190            }
191        }
192        Ok(hits)
193    }
194
195    /// **Vector⋈graph** (ADR 0011): rank the top-`k` candidates by semantic
196    /// similarity, then keep only those whose record falls in `root`'s
197    /// call-graph neighborhood in `graph` — `root` itself, its callers, and its
198    /// callees. Composition is by shared key: a record is in the neighborhood
199    /// when its [`RecordKey`](gonzalo_core::RecordKey)`.id` equals a symbol name
200    /// in the neighborhood (no new cross-store linkage, ADR 0008/0011).
201    ///
202    /// So "semantically similar to X **and** structurally near `root`" is one
203    /// call. The result is the in-neighborhood subset of the top-`k` (≤ `k`).
204    #[cfg(feature = "graph")]
205    pub async fn query_in_subgraph(
206        &self,
207        text: &str,
208        k: usize,
209        graph: &dyn gonzalo_graph::GraphStore,
210        root: &str,
211    ) -> Result<Vec<Hit>> {
212        use std::collections::BTreeSet;
213        let mut neighborhood: BTreeSet<String> = BTreeSet::from([root.to_string()]);
214        neighborhood.extend(graph.callers_of(root));
215        neighborhood.extend(graph.callees(root));
216
217        let hits = self.query(text, k, &KeyPrefix::default()).await?;
218        Ok(hits
219            .into_iter()
220            .filter(|h| neighborhood.contains(&h.record.key.id))
221            .collect())
222    }
223}
224
225/// The embeddable text for a record, or `None` if its kind is not
226/// knowledge-bearing (ADR 0011). Kept as the one-document view: the [`chunk`]s
227/// joined back into a single string.
228pub fn knowledge_text(record: &Record) -> Option<String> {
229    // A body that fails to parse is not a "one document"; collapse it to `None`
230    // for this convenience view. The indexing path ([`chunk`]/[`ingest`])
231    // surfaces the error instead (#139).
232    chunk(record).ok().flatten().map(|cs| cs.join("\n"))
233}
234
235/// Split a record into ordered chunk texts. Returns `Ok(None)` if the kind is
236/// definitionally not knowledge-bearing (`Checkpoint`/`GraphManifest`, ADR
237/// 0011), and `Err` if a knowledge-bearing body fails to parse — so a corrupt
238/// record is distinguishable from a genuinely non-indexable one (#139).
239/// Chunking is per-kind: a `Session` by turn, a long `MemoryTier`/`Ticket` by
240/// section/paragraph, a `Topic` by bullet. The kind's title/name is folded into
241/// the first chunk so it stays searchable. A kind that yields a single piece is
242/// the one-document fallback — identical to phase-1 behavior. Extraction goes
243/// through the `gonzalo-domain` typed views.
244pub fn chunk(record: &Record) -> Result<Option<Vec<String>>> {
245    let chunks = match record.kind {
246        RecordKind::MemoryTier => {
247            let t = MemoryTier::from_body(&record.body)?;
248            prepend_header(&t.name, split_paragraphs(&t.content))
249        }
250        RecordKind::Topic => {
251            let t = Topic::from_body(&record.body)?;
252            prepend_header(&t.slug, t.bullets)
253        }
254        RecordKind::Session => {
255            let s = Session::from_body(&record.body)?;
256            prepend_header(&s.name, s.turns.into_iter().map(|turn| turn.text).collect())
257        }
258        RecordKind::Ticket => {
259            let t = Ticket::from_body(&record.body)?;
260            let mut cs = vec![format!("{}\n{}", t.title, t.labels.join(" "))];
261            cs.extend(split_paragraphs(&t.body.markdown));
262            cs
263        }
264        RecordKind::TicketEvent => {
265            let e = TicketEvent::from_body(&record.body)?;
266            vec![e.body]
267        }
268        // Not knowledge-bearing: a checkpoint is opaque state; a graph manifest
269        // is a path -> content-hash map, not natural-language text (ADR 0011/0012).
270        // These are the ONLY definitionally-opaque kinds — a parse failure on a
271        // knowledge-bearing kind above propagates as an error rather than being
272        // silently indistinguishable from "not indexable" (#139).
273        RecordKind::Checkpoint | RecordKind::GraphManifest => return Ok(None),
274    };
275    Ok(Some(chunks))
276}
277
278/// Fold `header` into the first `pieces` chunk so it stays searchable; if
279/// `pieces` is empty, the header stands alone as the single chunk.
280fn prepend_header(header: &str, mut pieces: Vec<String>) -> Vec<String> {
281    match pieces.first_mut() {
282        Some(first) => *first = format!("{header}\n{first}"),
283        None => pieces.push(header.to_string()),
284    }
285    pieces
286}
287
288/// How many chunk matches to over-fetch per requested record in [`query`], so
289/// de-duping chunks back to their parents still yields up to `k` distinct
290/// records.
291const OVERFETCH: usize = 8;
292
293/// Separates a parent id from a chunk ordinal in a derived index key. ASCII Unit
294/// Separator (never appears in real ids, which may themselves contain `#`/`/`),
295/// so [`parent_key`] recovers the parent unambiguously.
296const CHUNK_SEP: char = '\u{1f}';
297
298/// The index key for chunk `ordinal` of the record at `parent`: same
299/// namespace/collection, with the ordinal folded into the `id` (so a
300/// [`KeyPrefix`] filter, which matches on namespace/collection only, is
301/// unaffected). Every chunk — including a single-chunk fallback — carries an
302/// ordinal, so keys are always parseable.
303fn chunk_key(parent: &gonzalo_core::RecordKey, ordinal: usize) -> gonzalo_core::RecordKey {
304    gonzalo_core::RecordKey::new(
305        &parent.namespace,
306        &parent.collection,
307        format!("{}{CHUNK_SEP}{ordinal}", parent.id),
308    )
309}
310
311/// Inverse of [`chunk_key`]: recover the parent key and chunk ordinal. A key
312/// without the separator (defensive) is treated as parent, ordinal 0.
313fn parent_key(chunk: &gonzalo_core::RecordKey) -> (gonzalo_core::RecordKey, usize) {
314    match chunk.id.rsplit_once(CHUNK_SEP) {
315        Some((id, ord)) => {
316            let ordinal = ord.parse().unwrap_or(0);
317            (
318                gonzalo_core::RecordKey::new(&chunk.namespace, &chunk.collection, id),
319                ordinal,
320            )
321        }
322        None => (chunk.clone(), 0),
323    }
324}
325
326/// Split text on blank-line boundaries into trimmed, non-empty paragraphs.
327fn split_paragraphs(text: &str) -> Vec<String> {
328    text.split("\n\n")
329        .map(str::trim)
330        .filter(|p| !p.is_empty())
331        .map(str::to_string)
332        .collect()
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338    use async_trait::async_trait;
339    use gonzalo_core::{Body, Identity, Meta, PutResult, RecordKey, Revision};
340    use gonzalo_domain::{BodyFormat, Provider, State, StateCategory, TicketBody};
341    use gonzalo_store_fs::FsStore;
342    use gonzalo_vector::MemoryVectorIndex;
343    use std::collections::BTreeMap;
344
345    /// A deterministic bag-of-words embedder: cosine similarity tracks word
346    /// overlap, enough to rank a matching record first.
347    struct Bow;
348    #[async_trait]
349    impl Embedder for Bow {
350        async fn embed(&self, text: &str) -> Result<Vec<f32>> {
351            let mut v = vec![0f32; 32];
352            for word in text.split_whitespace() {
353                let h = word.bytes().map(|b| b as usize).sum::<usize>() % 32;
354                v[h] += 1.0;
355            }
356            Ok(v)
357        }
358    }
359
360    fn record(key: &RecordKey, kind: RecordKind, body: Body) -> Record {
361        Record {
362            revision: Revision::initial(body.bytes()),
363            parent: None,
364            body,
365            kind,
366            meta: Meta {
367                author: Identity::new("t"),
368                origin_system: "test".into(),
369                created: 0,
370                updated: 0,
371                labels: BTreeMap::new(),
372            },
373            links: Vec::new(),
374            key: key.clone(),
375        }
376    }
377
378    #[test]
379    fn chunk_splits_session_by_turn() {
380        use gonzalo_domain::Turn;
381        let session = Session {
382            name: "sess".into(),
383            turns: vec![
384                Turn {
385                    role: "user".into(),
386                    text: "first turn about rust".into(),
387                },
388                Turn {
389                    role: "assistant".into(),
390                    text: "second turn about cooking".into(),
391                },
392                Turn {
393                    role: "user".into(),
394                    text: "third turn about music".into(),
395                },
396            ],
397        };
398        let key = RecordKey::new("caliban", "sessions", "s1");
399        let chunks = chunk(&record(
400            &key,
401            RecordKind::Session,
402            session.to_body().unwrap(),
403        ))
404        .unwrap()
405        .unwrap();
406        assert_eq!(chunks.len(), 3);
407        // Session name is prepended to the first chunk so it stays searchable.
408        assert!(chunks[0].contains("sess"));
409        assert!(chunks[0].contains("first turn about rust"));
410        assert!(chunks[1].contains("second turn about cooking"));
411        assert!(chunks[2].contains("third turn about music"));
412    }
413
414    #[test]
415    fn chunk_splits_topic_by_bullet_and_falls_back_to_one() {
416        // Multiple bullets -> one chunk each.
417        let topic = Topic {
418            slug: "rust".into(),
419            bullets: vec!["use clippy".into(), "run rustfmt".into()],
420        };
421        let key = RecordKey::new("caliban", "topics", "rust");
422        let chunks = chunk(&record(&key, RecordKind::Topic, topic.to_body().unwrap()))
423            .unwrap()
424            .unwrap();
425        assert_eq!(chunks.len(), 2);
426        assert!(chunks[0].contains("rust"));
427        assert!(chunks[0].contains("use clippy"));
428        assert!(chunks[1].contains("run rustfmt"));
429
430        // A single-piece record is the one-document fallback.
431        let one = Topic {
432            slug: "solo".into(),
433            bullets: vec!["only bullet".into()],
434        };
435        let chunks = chunk(&record(&key, RecordKind::Topic, one.to_body().unwrap()))
436            .unwrap()
437            .unwrap();
438        assert_eq!(chunks.len(), 1);
439
440        // Non-knowledge kinds still gate to None.
441        let ck = record(&key, RecordKind::Checkpoint, Body::Inline(b"{}".to_vec()));
442        assert_eq!(chunk(&ck).unwrap(), None);
443    }
444
445    #[test]
446    fn knowledge_text_per_kind() {
447        let topic = Topic {
448            slug: "rust".into(),
449            bullets: vec!["use clippy".into()],
450        };
451        let key = RecordKey::new("caliban", "topics", "rust");
452        let text = knowledge_text(&record(&key, RecordKind::Topic, topic.to_body().unwrap()));
453        assert_eq!(text.as_deref(), Some("rust\nuse clippy"));
454
455        // Checkpoint is not knowledge-bearing.
456        let ck = record(&key, RecordKind::Checkpoint, Body::Inline(b"{}".to_vec()));
457        assert_eq!(knowledge_text(&ck), None);
458    }
459
460    #[test]
461    fn ticket_text_includes_title_body_labels() {
462        let t = Ticket {
463            provider: Provider::GitHub,
464            uid: "o/r#1".into(),
465            display: "#1".into(),
466            item_type: "issue".into(),
467            title: "fix the parser".into(),
468            state: State {
469                category: StateCategory::Open,
470                resolution: None,
471                raw_name: "open".into(),
472                raw_id: None,
473            },
474            priority: None,
475            actors: vec![],
476            labels: vec!["bug".into()],
477            containers: vec![],
478            links: vec![],
479            body: TicketBody {
480                markdown: "the parser panics".into(),
481                format: BodyFormat::Markdown,
482                raw: None,
483            },
484            fields: BTreeMap::new(),
485        };
486        let key = RecordKey::new("tickets", "github", "o/r#1");
487        let text = knowledge_text(&record(&key, RecordKind::Ticket, t.to_body().unwrap())).unwrap();
488        assert!(text.contains("fix the parser"));
489        assert!(text.contains("the parser panics"));
490        assert!(text.contains("bug"));
491    }
492
493    async fn put(store: &FsStore, rec: Record) {
494        assert!(matches!(
495            store.put(rec, None).await.unwrap(),
496            PutResult::Committed(_)
497        ));
498    }
499
500    #[tokio::test]
501    async fn ingest_then_query_returns_matching_record() {
502        let dir = tempfile::tempdir().unwrap();
503        let store = FsStore::new(dir.path());
504
505        let rust = RecordKey::new("caliban", "topics", "rust");
506        let cooking = RecordKey::new("caliban", "topics", "cooking");
507        put(
508            &store,
509            record(
510                &rust,
511                RecordKind::Topic,
512                Topic {
513                    slug: "rust".into(),
514                    bullets: vec!["use clippy and cargo".into()],
515                }
516                .to_body()
517                .unwrap(),
518            ),
519        )
520        .await;
521        put(
522            &store,
523            record(
524                &cooking,
525                RecordKind::Topic,
526                Topic {
527                    slug: "cooking".into(),
528                    bullets: vec!["simmer the sauce slowly".into()],
529                }
530                .to_body()
531                .unwrap(),
532            ),
533        )
534        .await;
535
536        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
537        assert!(ks.ingest(&rust).await.unwrap());
538        assert!(ks.ingest(&cooking).await.unwrap());
539
540        let hits = ks
541            .query("clippy cargo", 1, &KeyPrefix::default())
542            .await
543            .unwrap();
544        assert_eq!(hits.len(), 1);
545        assert_eq!(hits[0].record.key, rust);
546    }
547
548    fn turn(text: &str) -> gonzalo_domain::Turn {
549        gonzalo_domain::Turn {
550            role: "user".into(),
551            text: text.into(),
552        }
553    }
554
555    #[tokio::test]
556    async fn query_ranks_a_matching_turn_over_a_diluted_decoy_and_dedups() {
557        let dir = tempfile::tempdir().unwrap();
558        let store = FsStore::new(dir.path());
559
560        // A session whose one-document average is heavily diluted: two short
561        // turns match "gamma", one long turn is all "alpha". Averaged into one
562        // vector, the alpha turn drowns out the gamma signal.
563        let sess = RecordKey::new("caliban", "sessions", "s1");
564        put(
565            &store,
566            record(
567                &sess,
568                RecordKind::Session,
569                Session {
570                    name: "s".into(),
571                    turns: vec![
572                        turn("gamma"),
573                        turn("gamma"),
574                        turn("alpha alpha alpha alpha alpha"),
575                    ],
576                }
577                .to_body()
578                .unwrap(),
579            ),
580        )
581        .await;
582
583        // A decoy whose whole (short) document is closer to "gamma" than the
584        // session's diluted average — so under one-document embedding the decoy
585        // outranks the session.
586        let decoy = RecordKey::new("caliban", "topics", "decoy");
587        put(
588            &store,
589            record(
590                &decoy,
591                RecordKind::Topic,
592                Topic {
593                    slug: "gamma".into(),
594                    bullets: vec!["alpha".into()],
595                }
596                .to_body()
597                .unwrap(),
598            ),
599        )
600        .await;
601
602        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
603        assert!(ks.ingest(&sess).await.unwrap());
604        assert!(ks.ingest(&decoy).await.unwrap());
605
606        let hits = ks.query("gamma", 5, &KeyPrefix::default()).await.unwrap();
607
608        // Chunking lets the matching turn score on its own, so the session ranks
609        // first...
610        assert_eq!(hits[0].record.key, sess, "matching turn should rank first");
611        // ...and the two matching turns de-dup to a single hit for the record.
612        assert_eq!(
613            hits.iter().filter(|h| h.record.key == sess).count(),
614            1,
615            "session should appear exactly once"
616        );
617    }
618
619    #[tokio::test]
620    async fn query_chunks_returns_individual_chunks_with_ordinals() {
621        let dir = tempfile::tempdir().unwrap();
622        let store = FsStore::new(dir.path());
623        let sess = RecordKey::new("caliban", "sessions", "s1");
624        put(
625            &store,
626            record(
627                &sess,
628                RecordKind::Session,
629                Session {
630                    name: "s".into(),
631                    turns: vec![turn("gamma"), turn("gamma"), turn("alpha")],
632                }
633                .to_body()
634                .unwrap(),
635            ),
636        )
637        .await;
638
639        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
640        assert!(ks.ingest(&sess).await.unwrap());
641
642        // The two "gamma" turns are the two closest chunks — returned
643        // individually, NOT de-duped to the record.
644        let hits = ks
645            .query_chunks("gamma", 2, &KeyPrefix::default())
646            .await
647            .unwrap();
648        assert_eq!(hits.len(), 2);
649        assert!(hits.iter().all(|h| h.record.key == sess));
650        let ordinals: std::collections::BTreeSet<usize> = hits.iter().map(|h| h.ordinal).collect();
651        assert_eq!(ordinals, std::collections::BTreeSet::from([0, 1]));
652        assert!(hits.iter().all(|h| h.text.contains("gamma")));
653    }
654
655    #[tokio::test]
656    async fn reingest_shrink_removes_orphan_chunks() {
657        let dir = tempfile::tempdir().unwrap();
658        let sess = RecordKey::new("caliban", "sessions", "s1");
659        let ks = KnowledgeStore::new(FsStore::new(dir.path()), MemoryVectorIndex::default(), Bow);
660
661        // v1: three turns -> three chunks.
662        put(
663            ks.store(),
664            record(
665                &sess,
666                RecordKind::Session,
667                Session {
668                    name: "s".into(),
669                    turns: vec![turn("gamma"), turn("delta"), turn("epsilon")],
670                }
671                .to_body()
672                .unwrap(),
673            ),
674        )
675        .await;
676        assert!(ks.ingest(&sess).await.unwrap());
677
678        // The third turn's content is indexed while it exists.
679        let before = ks
680            .query_chunks("epsilon", 10, &KeyPrefix::default())
681            .await
682            .unwrap();
683        assert!(before.iter().any(|h| h.text.contains("epsilon")));
684
685        // v2: shrink to a single turn. The record now has one chunk; the "delta"
686        // and "epsilon" chunks are orphaned in the index.
687        let current = ks.store().get(&sess).await.unwrap().unwrap();
688        let body = Session {
689            name: "s".into(),
690            turns: vec![turn("gamma")],
691        }
692        .to_body()
693        .unwrap();
694        let mut v2 = record(&sess, RecordKind::Session, body.clone());
695        v2.parent = Some(current.revision.clone());
696        v2.revision = current.revision.next(body.bytes());
697        assert!(matches!(
698            ks.store().put(v2, Some(current.revision)).await.unwrap(),
699            PutResult::Committed(_)
700        ));
701        assert!(ks.ingest(&sess).await.unwrap());
702
703        // The orphaned "epsilon" (and "delta") chunks are gone from the index;
704        // only the surviving "gamma" chunk remains.
705        let after = ks
706            .query_chunks("epsilon", 10, &KeyPrefix::default())
707            .await
708            .unwrap();
709        assert!(
710            !after.iter().any(|h| h.text.contains("epsilon")),
711            "orphaned chunk should have been removed on re-ingest"
712        );
713        assert!(
714            !after.iter().any(|h| h.text.contains("delta")),
715            "orphaned chunk should have been removed on re-ingest"
716        );
717        assert!(
718            after.iter().any(|h| h.text.contains("gamma")),
719            "the surviving turn should still be indexed"
720        );
721    }
722
723    #[tokio::test]
724    async fn ingest_skips_non_knowledge_kinds() {
725        let dir = tempfile::tempdir().unwrap();
726        let store = FsStore::new(dir.path());
727        let key = RecordKey::new("caliban", "checkpoints", "c1");
728        put(
729            &store,
730            record(&key, RecordKind::Checkpoint, Body::Inline(b"{}".to_vec())),
731        )
732        .await;
733        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
734        assert!(!ks.ingest(&key).await.unwrap(), "checkpoint is not indexed");
735    }
736
737    #[tokio::test]
738    async fn ingest_surfaces_body_parse_error_and_gates_opaque_kinds() {
739        let dir = tempfile::tempdir().unwrap();
740        let store = FsStore::new(dir.path());
741
742        // A knowledge-bearing kind (Topic) whose body bytes are not a valid
743        // Topic: a corrupt record is a real failure, not "not indexable" (#139).
744        let corrupt = RecordKey::new("caliban", "topics", "corrupt");
745        put(
746            &store,
747            record(
748                &corrupt,
749                RecordKind::Topic,
750                Body::Inline(b"not a valid topic".to_vec()),
751            ),
752        )
753        .await;
754
755        // A definitionally-opaque kind still gates cleanly to Ok(false).
756        let ck = RecordKey::new("caliban", "checkpoints", "c1");
757        put(
758            &store,
759            record(
760                &ck,
761                RecordKind::Checkpoint,
762                Body::Inline(b"opaque".to_vec()),
763            ),
764        )
765        .await;
766
767        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
768        assert!(
769            ks.ingest(&corrupt).await.is_err(),
770            "a corrupt knowledge-bearing body must surface an error, not Ok(false)"
771        );
772        assert!(
773            !ks.ingest(&ck).await.unwrap(),
774            "a genuinely-opaque kind is non-indexable: Ok(false)"
775        );
776    }
777
778    #[tokio::test]
779    async fn remove_deindexes_a_records_chunk_vectors() {
780        let dir = tempfile::tempdir().unwrap();
781        let store = FsStore::new(dir.path());
782
783        // Three topics all matching "gamma". One will be removed; the other two
784        // are the live set.
785        let removed = RecordKey::new("caliban", "topics", "removed");
786        let live1 = RecordKey::new("caliban", "topics", "live1");
787        let live2 = RecordKey::new("caliban", "topics", "live2");
788        for key in [&removed, &live1, &live2] {
789            put(
790                &store,
791                record(
792                    key,
793                    RecordKind::Topic,
794                    Topic {
795                        slug: "gamma".into(),
796                        bullets: vec!["gamma".into()],
797                    }
798                    .to_body()
799                    .unwrap(),
800                ),
801            )
802            .await;
803        }
804
805        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
806        for key in [&removed, &live1, &live2] {
807            assert!(ks.ingest(key).await.unwrap());
808        }
809
810        // De-index the one record. Its chunk vectors must no longer match, so
811        // the k=2 query is filled entirely by the live records — no shrinkage
812        // from orphaned vectors, and the removed record never appears (#150).
813        ks.remove(&removed).await.unwrap();
814
815        let hits = ks.query("gamma", 2, &KeyPrefix::default()).await.unwrap();
816        assert_eq!(hits.len(), 2, "k=2 should be filled by the live records");
817        let keys: std::collections::BTreeSet<RecordKey> =
818            hits.iter().map(|h| h.record.key.clone()).collect();
819        assert!(
820            !keys.contains(&removed),
821            "removed record must not appear after de-indexing"
822        );
823        assert_eq!(
824            keys,
825            std::collections::BTreeSet::from([live1, live2]),
826            "only the live records should remain"
827        );
828    }
829
830    #[cfg(feature = "graph")]
831    #[tokio::test]
832    async fn query_in_subgraph_restricts_to_the_neighborhood() {
833        use gonzalo_graph::{GraphStore, InMemoryGraphStore, build_rust};
834
835        let dir = tempfile::tempdir().unwrap();
836        let store = FsStore::new(dir.path());
837
838        // Three topic records keyed by symbol name, all mentioning "database",
839        // so all three are semantically similar to the query.
840        for id in ["root", "helper", "faraway"] {
841            let key = RecordKey::new("code", "symbols", id);
842            let topic = Topic {
843                slug: id.into(),
844                bullets: vec!["touches the database layer".into()],
845            };
846            put(
847                &store,
848                record(&key, RecordKind::Topic, topic.to_body().unwrap()),
849            )
850            .await;
851        }
852        let ks = KnowledgeStore::new(store, MemoryVectorIndex::default(), Bow);
853        for id in ["root", "helper", "faraway"] {
854            assert!(
855                ks.ingest(&RecordKey::new("code", "symbols", id))
856                    .await
857                    .unwrap()
858            );
859        }
860
861        // Call graph: root -> helper; faraway is unrelated. So root's
862        // neighborhood is {root, helper}.
863        let mut graph = InMemoryGraphStore::new();
864        graph.insert(
865            "lib.rs",
866            build_rust("fn root() { helper(); }\nfn helper() {}\nfn faraway() {}"),
867        );
868
869        let hits = ks
870            .query_in_subgraph("database", 10, &graph, "root")
871            .await
872            .unwrap();
873        let ids: BTreeMap<String, ()> =
874            hits.iter().map(|h| (h.record.key.id.clone(), ())).collect();
875        assert!(ids.contains_key("root"));
876        assert!(ids.contains_key("helper"));
877        assert!(
878            !ids.contains_key("faraway"),
879            "faraway is semantically similar but outside root's neighborhood"
880        );
881    }
882}