1use gonzalo_core::{KeyPrefix, Record, RecordKind, Result, Store};
19use gonzalo_domain::{MemoryTier, RecordCodec, Session, Ticket, TicketEvent, Topic};
20use gonzalo_vector::{Embedder, VectorIndex};
21
22#[derive(Debug, Clone)]
25pub struct Hit {
26 pub record: Record,
27 pub score: f32,
28}
29
30#[derive(Debug, Clone)]
34pub struct ChunkHit {
35 pub record: Record,
36 pub score: f32,
37 pub ordinal: usize,
38 pub text: String,
39}
40
41pub struct KnowledgeStore<S, V, E> {
44 store: S,
45 index: V,
46 embedder: E,
47 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 pub fn store(&self) -> &S {
66 &self.store
67 }
68
69 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 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 pub async fn remove(&self, key: &gonzalo_core::RecordKey) -> Result<()> {
113 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 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 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 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 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 #[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
225pub fn knowledge_text(record: &Record) -> Option<String> {
229 chunk(record).ok().flatten().map(|cs| cs.join("\n"))
233}
234
235pub 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 RecordKind::Checkpoint | RecordKind::GraphManifest => return Ok(None),
274 };
275 Ok(Some(chunks))
276}
277
278fn 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
288const OVERFETCH: usize = 8;
292
293const CHUNK_SEP: char = '\u{1f}';
297
298fn 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
311fn 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
326fn 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 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 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 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 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 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 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 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 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 assert_eq!(hits[0].record.key, sess, "matching turn should rank first");
611 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 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 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 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 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 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 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 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 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 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 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 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}