1#![allow(missing_docs)]
10
11use brain_db_sdk::RecallAnswer;
12use brain_db_sdk::wire::types::{
13 AnswerKindWire, EdgeDirectionWire, EncodeGraphEdge, EncodeGraphNode, EncodeResponse,
14 EncodeStageArtifact, EncodeStageGraph, EncodeStageKeywordField, EncodeStageRecord,
15 ForgetResponse, MemoryEdgeWire, MemoryInspectResponse, MemoryListDirWire, MemoryListItem,
16 MemoryListRequest, MemoryListSortWire, MemoryListTimeAxisWire, MemoryResult, WireMemoryId,
17};
18use serde::{Deserialize, Serialize};
19
20use super::{hex_decode, hex_encode, mem_id_decimal, parse_memory_id};
21
22#[derive(Debug, Deserialize)]
24pub struct EncodeBody {
25 pub text: String,
27 #[serde(default)]
29 pub session: Option<u64>,
30 #[serde(default)]
32 pub occurred_at: Option<u64>,
33}
34
35#[derive(Debug, Serialize)]
37pub struct EncodeResponseDto {
38 pub memory_id: String,
39 pub was_deduplicated: bool,
40 pub salience: f32,
41 pub kind: u8,
42 pub created_at_unix_nanos: u64,
43 pub auto_edges_added: u32,
44}
45
46impl From<EncodeResponse> for EncodeResponseDto {
47 fn from(r: EncodeResponse) -> Self {
48 Self {
49 memory_id: r.memory_id.to_string(),
50 was_deduplicated: r.was_deduplicated,
51 salience: r.salience,
52 kind: r.kind as u8,
53 created_at_unix_nanos: r.created_at_unix_nanos,
54 auto_edges_added: r.auto_edges_added,
55 }
56 }
57}
58
59#[derive(Debug, Deserialize)]
61pub struct RecallBody {
62 #[serde(alias = "cue")]
63 pub query: String,
64 #[serde(default)]
65 pub max_results: Option<u32>,
66 #[serde(default)]
67 pub subject: Option<String>,
68 #[serde(default)]
72 pub scope: Option<String>,
73}
74
75#[derive(Debug, Serialize)]
77pub struct MemoryHit {
78 pub memory_id: String,
79 pub text: String,
80 pub similarity_score: f32,
81 pub confidence: f32,
82 pub salience: f32,
83 pub kind: u8,
84 pub created_at_unix_nanos: u64,
85}
86
87impl From<MemoryResult> for MemoryHit {
88 fn from(m: MemoryResult) -> Self {
89 Self {
90 memory_id: m.memory_id.to_string(),
91 text: m.text,
92 similarity_score: m.similarity_score,
93 confidence: m.confidence,
94 salience: m.salience,
95 kind: m.kind as u8,
96 created_at_unix_nanos: m.created_at_unix_nanos,
97 }
98 }
99}
100
101#[derive(Debug, Serialize)]
103pub struct RecallResponseDto {
104 pub answer_kind: &'static str,
106 pub memories: Vec<MemoryHit>,
107}
108
109impl From<RecallAnswer> for RecallResponseDto {
110 fn from(a: RecallAnswer) -> Self {
111 let answer_kind = match a.answer_kind {
112 AnswerKindWire::Single => "single",
113 AnswerKindWire::Many => "many",
114 AnswerKindWire::None => "none",
115 };
116 Self {
117 answer_kind,
118 memories: a.memories.into_iter().map(MemoryHit::from).collect(),
119 }
120 }
121}
122
123#[derive(Debug, Deserialize)]
125pub struct ForgetBody {
126 pub memory_id: String,
128 #[serde(default)]
129 pub hard: bool,
130}
131
132impl ForgetBody {
133 pub fn parse_id(&self) -> Result<WireMemoryId, String> {
135 parse_memory_id(&self.memory_id)
136 }
137}
138
139#[derive(Debug, Serialize)]
141pub struct ForgetResponseDto {
142 pub memory_id: String,
143 pub was_already_forgotten: bool,
144 pub edges_removed: u32,
145}
146
147impl From<ForgetResponse> for ForgetResponseDto {
148 fn from(r: ForgetResponse) -> Self {
149 Self {
150 memory_id: r.memory_id.to_string(),
151 was_already_forgotten: r.was_already_forgotten,
152 edges_removed: r.edges_removed,
153 }
154 }
155}
156
157const LIST_DEFAULT_LIMIT: u32 = 50;
161const LIST_MAX_LIMIT: u32 = 100;
162
163#[derive(Debug, Deserialize)]
168pub struct MemoryListQuery {
169 #[serde(default)]
171 pub limit: Option<u32>,
172 #[serde(default)]
174 pub cursor: Option<String>,
175 #[serde(default)]
177 pub dir: Option<String>,
178 #[serde(default)]
180 pub include_tombstoned: Option<bool>,
181}
182
183impl MemoryListQuery {
184 pub fn to_request(&self) -> Result<MemoryListRequest, String> {
186 let limit = self
187 .limit
188 .unwrap_or(LIST_DEFAULT_LIMIT)
189 .clamp(1, LIST_MAX_LIMIT);
190 let dir = match self.dir.as_deref().map(str::trim) {
191 None | Some("" | "desc") => MemoryListDirWire::Desc,
192 Some("asc") => MemoryListDirWire::Asc,
193 Some(other) => return Err(format!("dir must be asc|desc, got `{other}`")),
194 };
195 let cursor = match &self.cursor {
196 Some(c) if !c.is_empty() => hex_decode(c)?,
197 _ => Vec::new(),
198 };
199 Ok(MemoryListRequest {
200 sort: MemoryListSortWire::Created,
201 dir,
202 limit,
203 cursor,
204 kinds: Vec::new(),
205 include_tombstoned: self.include_tombstoned.unwrap_or(false),
206 time_axis: MemoryListTimeAxisWire::Created,
207 from_unix_nanos: 0,
208 to_unix_nanos: 0,
209 salience_min: 0.0,
210 salience_max: 1.0,
211 text_contains: String::new(),
212 act_as: None,
213 })
214 }
215}
216
217#[derive(Debug, Serialize)]
221pub struct MemoryListItemDto {
222 pub memory_id: String,
224 pub text: String,
225 pub kind: u8,
227 pub state: u8,
229 pub created_at_unix_nanos: u64,
230 pub occurred_at_unix_nanos: u64,
231 pub last_accessed_at_unix_nanos: u64,
232 pub salience: f32,
233 pub access_count: u32,
234 pub statement_count: u32,
235 pub entity_count: u32,
236 pub relation_count: u32,
237}
238
239impl From<MemoryListItem> for MemoryListItemDto {
240 fn from(m: MemoryListItem) -> Self {
241 Self {
242 memory_id: mem_id_decimal(&m.memory_id),
243 text: m.text,
244 kind: m.kind,
245 state: m.state,
246 created_at_unix_nanos: m.created_at_unix_nanos,
247 occurred_at_unix_nanos: m.occurred_at_unix_nanos,
248 last_accessed_at_unix_nanos: m.last_accessed_at_unix_nanos,
249 salience: m.salience,
250 access_count: m.access_count,
251 statement_count: m.statement_count,
252 entity_count: m.entity_count,
253 relation_count: m.relation_count,
254 }
255 }
256}
257
258#[derive(Debug, Serialize)]
261pub struct MemoryListPageDto {
262 pub items: Vec<MemoryListItemDto>,
263 #[serde(skip_serializing_if = "Option::is_none")]
264 pub next_cursor: Option<String>,
265}
266
267impl MemoryListPageDto {
268 pub fn from_frames(frames: Vec<brain_db_sdk::wire::types::MemoryListResponseFrame>) -> Self {
271 let next = frames
272 .last()
273 .map(|f| f.next_cursor.clone())
274 .filter(|c| !c.is_empty())
275 .map(|c| hex_encode(&c));
276 let items = frames
277 .into_iter()
278 .flat_map(|f| f.items)
279 .map(MemoryListItemDto::from)
280 .collect();
281 Self {
282 items,
283 next_cursor: next,
284 }
285 }
286}
287
288#[derive(Debug, Serialize)]
297pub struct MemoryInspectDto {
298 pub found: bool,
299 pub memory_id: String,
301 pub text: String,
302 pub artifact: StageArtifactDto,
303 pub edges: Vec<MemoryEdgeDto>,
310}
311
312#[derive(Debug, Serialize)]
317pub struct MemoryEdgeDto {
318 pub kind: &'static str,
320 pub direction: &'static str,
322 pub peer: String,
324 pub weight: f32,
325 pub created_at_unix_nanos: u64,
326 #[serde(skip_serializing_if = "Option::is_none")]
327 pub annotation: Option<String>,
328}
329
330impl From<MemoryEdgeWire> for MemoryEdgeDto {
331 fn from(e: MemoryEdgeWire) -> Self {
332 Self {
333 kind: super::graph::edge_kind_str(e.kind),
334 direction: match e.direction {
335 EdgeDirectionWire::Outgoing => "outgoing",
336 EdgeDirectionWire::Incoming => "incoming",
337 },
338 peer: e.peer.to_string(),
339 weight: e.weight,
340 created_at_unix_nanos: e.created_at_unix_nanos,
341 annotation: e.annotation,
342 }
343 }
344}
345
346impl From<MemoryInspectResponse> for MemoryInspectDto {
347 fn from(r: MemoryInspectResponse) -> Self {
348 Self {
349 found: r.found,
350 memory_id: mem_id_decimal(&r.memory_id),
351 text: r.text,
352 artifact: StageArtifactDto::from(r.artifact),
353 edges: r.edges.into_iter().map(MemoryEdgeDto::from).collect(),
354 }
355 }
356}
357
358#[derive(Debug, Serialize)]
361pub struct StageArtifactDto {
362 pub vector: Vec<f32>,
364 pub record: Option<StageRecordDto>,
366 pub hype_questions: Vec<String>,
368 pub keyword_fields: Vec<StageKeywordFieldDto>,
370 pub graph: Option<StageGraphDto>,
372}
373
374impl From<EncodeStageArtifact> for StageArtifactDto {
375 fn from(a: EncodeStageArtifact) -> Self {
376 Self {
377 vector: a.vector,
378 record: a.record.map(StageRecordDto::from),
379 hype_questions: a.hype_questions,
380 keyword_fields: a
381 .keyword_fields
382 .into_iter()
383 .map(StageKeywordFieldDto::from)
384 .collect(),
385 graph: a.graph.map(StageGraphDto::from),
386 }
387 }
388}
389
390#[derive(Debug, Serialize)]
392pub struct StageRecordDto {
393 pub memory_id: String,
394 pub kind: u8,
395 pub salience: f32,
396 pub created_at_unix_nanos: u64,
397 pub occurred_at_unix_nanos: u64,
398 pub vector_dim: u32,
399 pub text_len: u32,
400 pub lsn: u64,
401}
402
403impl From<EncodeStageRecord> for StageRecordDto {
404 fn from(r: EncodeStageRecord) -> Self {
405 Self {
406 memory_id: mem_id_decimal(&r.memory_id),
407 kind: r.kind,
408 salience: r.salience,
409 created_at_unix_nanos: r.created_at_unix_nanos,
410 occurred_at_unix_nanos: r.occurred_at_unix_nanos,
411 vector_dim: r.vector_dim,
412 text_len: r.text_len,
413 lsn: r.lsn,
414 }
415 }
416}
417
418#[derive(Debug, Serialize)]
420pub struct StageKeywordFieldDto {
421 pub field: String,
422 pub terms: Vec<String>,
423}
424
425impl From<EncodeStageKeywordField> for StageKeywordFieldDto {
426 fn from(k: EncodeStageKeywordField) -> Self {
427 Self {
428 field: k.field,
429 terms: k.terms,
430 }
431 }
432}
433
434#[derive(Debug, Serialize)]
436pub struct StageGraphDto {
437 pub nodes: Vec<StageGraphNodeDto>,
438 pub edges: Vec<StageGraphEdgeDto>,
439}
440
441impl From<EncodeStageGraph> for StageGraphDto {
442 fn from(g: EncodeStageGraph) -> Self {
443 Self {
444 nodes: g.nodes.into_iter().map(StageGraphNodeDto::from).collect(),
445 edges: g.edges.into_iter().map(StageGraphEdgeDto::from).collect(),
446 }
447 }
448}
449
450#[derive(Debug, Serialize)]
451pub struct StageGraphNodeDto {
452 pub id: String,
454 pub name: String,
455 pub kind: String,
456 pub type_qname: String,
457}
458
459impl From<EncodeGraphNode> for StageGraphNodeDto {
460 fn from(n: EncodeGraphNode) -> Self {
461 Self {
462 id: hex_encode(&n.id),
463 name: n.name,
464 kind: n.kind,
465 type_qname: n.type_qname,
466 }
467 }
468}
469
470#[derive(Debug, Serialize)]
471pub struct StageGraphEdgeDto {
472 pub source: String,
474 pub target: String,
476 pub predicate: String,
477 pub kind: String,
478 pub confidence: f32,
479 #[serde(skip_serializing_if = "Option::is_none")]
483 pub event_at_unix_nanos: Option<u64>,
484}
485
486impl From<EncodeGraphEdge> for StageGraphEdgeDto {
487 fn from(e: EncodeGraphEdge) -> Self {
488 Self {
489 source: hex_encode(&e.source),
490 target: hex_encode(&e.target),
491 predicate: e.predicate,
492 kind: e.kind,
493 confidence: e.confidence,
494 event_at_unix_nanos: e.event_at_unix_nanos,
495 }
496 }
497}
498
499#[cfg(test)]
500mod memory_tests {
501 use super::*;
502 use brain_db_sdk::RecallAnswer;
503 use brain_db_sdk::wire::types::{
504 EdgeKindWire, EncodeGraphEdge, EncodeGraphNode, EncodeResponse, EncodeStageArtifact,
505 EncodeStageGraph, EncodeStageKeywordField, EncodeStageRecord, ForgetResponse,
506 MemoryInspectResponse, MemoryKindWire, MemoryListItem, MemoryListResponseFrame,
507 MemoryResult,
508 };
509
510 fn sample_result(tag: u8, text: &str) -> MemoryResult {
511 MemoryResult {
512 memory_id: u128::from(tag),
513 text: text.into(),
514 similarity_score: 0.9,
515 confidence: 0.8,
516 salience: 0.5,
517 kind: MemoryKindWire::Semantic,
518 space_id: [tag; 16],
519 session_id: 0,
520 created_at_unix_nanos: 1,
521 last_accessed_at_unix_nanos: 1,
522 edges: None,
523 contributing_retrievers: vec![],
524 fused_score: 0.7,
525 rerank_score: None,
526 salience_initial: 0.5,
527 access_count: 0,
528 lsn: 1,
529 flags: 0,
530 consolidated_at_unix_nanos: None,
531 occurred_at_unix_nanos: None,
532 edges_out_count: 0,
533 edges_in_count: 0,
534 graph: None,
535 }
536 }
537
538 #[test]
539 fn list_query_clamps_limit_defaults_dir_and_rejects_bad_dir() {
540 let over = MemoryListQuery {
542 limit: Some(9999),
543 cursor: None,
544 dir: None,
545 include_tombstoned: None,
546 };
547 let req = over.to_request().unwrap();
548 assert_eq!(req.limit, LIST_MAX_LIMIT);
549 assert!(matches!(req.dir, MemoryListDirWire::Desc)); assert_eq!(req.cursor, [] as [u8; 0]);
551 assert!(!req.include_tombstoned);
552
553 let asc = MemoryListQuery {
554 limit: None,
555 cursor: Some("2a00ff".into()),
556 dir: Some(" asc ".into()),
557 include_tombstoned: Some(true),
558 }
559 .to_request()
560 .unwrap();
561 assert_eq!(asc.limit, LIST_DEFAULT_LIMIT);
562 assert!(matches!(asc.dir, MemoryListDirWire::Asc));
563 assert_eq!(asc.cursor, vec![0x2a, 0x00, 0xff]); assert!(asc.include_tombstoned);
565
566 assert!(
567 MemoryListQuery {
568 limit: None,
569 cursor: None,
570 dir: Some("sideways".into()),
571 include_tombstoned: None,
572 }
573 .to_request()
574 .is_err()
575 );
576 assert!(
578 MemoryListQuery {
579 limit: None,
580 cursor: Some("zz".into()),
581 dir: None,
582 include_tombstoned: None,
583 }
584 .to_request()
585 .is_err()
586 );
587 }
588
589 #[test]
590 fn forget_body_parses_decimal_id_and_rejects_junk() {
591 let ok = ForgetBody {
592 memory_id: "42".into(),
593 hard: true,
594 };
595 assert_eq!(ok.parse_id().unwrap(), 42u128);
596 assert!(
597 ForgetBody {
598 memory_id: "not-a-number".into(),
599 hard: false,
600 }
601 .parse_id()
602 .is_err()
603 );
604 }
605
606 #[test]
607 fn encode_response_maps_every_field_and_kind_byte() {
608 let resp = EncodeResponse {
609 memory_id: 7,
610 was_deduplicated: true,
611 salience: 0.42,
612 auto_edges_added: 3,
613 lsn: 99,
614 space_id: [0xAB; 16],
615 session_id: 5,
616 kind: MemoryKindWire::Consolidated,
617 created_at_unix_nanos: 1_700_000_000,
618 edges_out_count: 0,
619 embedding_model_fp: [0; 16],
620 pending_stages: vec![],
621 has_active_schema: true,
622 trace: None,
623 };
624 let dto = EncodeResponseDto::from(resp);
625 assert_eq!(dto.memory_id, "7"); assert!(dto.was_deduplicated);
627 assert_eq!(dto.salience.to_bits(), 0.42f32.to_bits());
628 assert_eq!(dto.kind, 2); assert_eq!(dto.created_at_unix_nanos, 1_700_000_000);
630 assert_eq!(dto.auto_edges_added, 3);
631 }
632
633 #[test]
634 fn recall_answer_kind_maps_single_many_none() {
635 let cases = [
636 (AnswerKindWire::Single, "single"),
637 (AnswerKindWire::Many, "many"),
638 (AnswerKindWire::None, "none"),
639 ];
640 for (kind, expected) in cases {
641 let answer = RecallAnswer {
642 answer_kind: kind,
643 memories: vec![sample_result(0xAA, "hit")],
644 };
645 let dto = RecallResponseDto::from(answer);
646 assert_eq!(dto.answer_kind, expected);
647 assert_eq!(dto.memories.len(), 1);
648 assert_eq!(dto.memories[0].memory_id, "170"); assert_eq!(dto.memories[0].text, "hit");
650 assert_eq!(dto.memories[0].kind, 1); }
652 }
653
654 #[test]
655 fn forget_response_maps_fields() {
656 let dto = ForgetResponseDto::from(ForgetResponse {
657 memory_id: 42,
658 was_already_forgotten: true,
659 edges_removed: 4,
660 });
661 assert_eq!(dto.memory_id, "42");
662 assert!(dto.was_already_forgotten);
663 assert_eq!(dto.edges_removed, 4);
664 }
665
666 fn sample_list_item(tag: u8) -> MemoryListItem {
667 MemoryListItem {
668 memory_id: [tag; 16],
669 space_id: [0; 16],
670 session_id: 0,
671 text: "row".into(),
672 kind: 1,
673 state: 0,
674 created_at_unix_nanos: 10,
675 occurred_at_unix_nanos: 0,
676 last_accessed_at_unix_nanos: 11,
677 salience: 0.5,
678 access_count: 2,
679 source_request_id: [0; 16],
680 statement_count: 1,
681 entity_count: 2,
682 relation_count: 3,
683 }
684 }
685
686 #[test]
687 fn list_page_flattens_frames_and_hex_encodes_tail_cursor() {
688 let frames = vec![
690 MemoryListResponseFrame {
691 items: vec![sample_list_item(1)],
692 next_cursor: Vec::new(),
693 cumulative_count: 1,
694 is_final: false,
695 },
696 MemoryListResponseFrame {
697 items: vec![sample_list_item(2)],
698 next_cursor: vec![0x2a, 0xff],
699 cumulative_count: 2,
700 is_final: true,
701 },
702 ];
703 let dto = MemoryListPageDto::from_frames(frames);
704 assert_eq!(dto.items.len(), 2);
705 assert_eq!(dto.items[0].statement_count, 1);
706 assert_eq!(dto.items[1].relation_count, 3);
707 assert_eq!(dto.next_cursor.as_deref(), Some("2aff"));
708
709 let done = MemoryListPageDto::from_frames(vec![MemoryListResponseFrame {
711 items: vec![sample_list_item(3)],
712 next_cursor: Vec::new(),
713 cumulative_count: 1,
714 is_final: true,
715 }]);
716 assert!(done.next_cursor.is_none());
717 }
718
719 #[test]
720 fn inspect_maps_stages_and_hex_encodes_graph_ids() {
721 let resp = MemoryInspectResponse {
722 found: true,
723 memory_id: [0; 16],
724 text: "hello".into(),
725 artifact: EncodeStageArtifact {
726 vector: vec![0.1, 0.2],
727 record: Some(EncodeStageRecord {
728 memory_id: [0; 16],
729 kind: 1,
730 salience: 0.5,
731 created_at_unix_nanos: 1,
732 occurred_at_unix_nanos: 0,
733 vector_dim: 384,
734 text_len: 5,
735 lsn: 7,
736 }),
737 hype_questions: vec!["who?".into()],
738 keyword_fields: vec![EncodeStageKeywordField {
739 field: "text".into(),
740 terms: vec!["hello".into()],
741 }],
742 graph: Some(EncodeStageGraph {
743 nodes: vec![EncodeGraphNode {
744 id: [0x11; 16],
745 name: "Ada".into(),
746 kind: "entity".into(),
747 type_qname: "brain:Person".into(),
748 }],
749 edges: vec![EncodeGraphEdge {
750 source: [0x11; 16],
751 target: [0x22; 16],
752 predicate: "born_in".into(),
753 kind: "statement".into(),
754 confidence: 0.9,
755 event_at_unix_nanos: Some(123),
756 }],
757 }),
758 },
759 edges: vec![
760 MemoryEdgeWire {
761 kind: EdgeKindWire::Caused,
762 direction: EdgeDirectionWire::Outgoing,
763 peer: 42,
764 weight: 0.75,
765 created_at_unix_nanos: 900,
766 annotation: Some("manually asserted".into()),
767 },
768 MemoryEdgeWire {
769 kind: EdgeKindWire::Supports,
770 direction: EdgeDirectionWire::Incoming,
771 peer: 7,
772 weight: 0.25,
773 created_at_unix_nanos: 901,
774 annotation: None,
775 },
776 ],
777 };
778 let dto = MemoryInspectDto::from(resp);
779 assert!(dto.found);
780 assert_eq!(dto.text, "hello");
781 assert_eq!(dto.artifact.vector, vec![0.1, 0.2]);
782 assert_eq!(dto.artifact.record.as_ref().unwrap().vector_dim, 384);
783 assert_eq!(dto.artifact.hype_questions, vec!["who?".to_string()]);
784 let graph = dto.artifact.graph.unwrap();
785 assert_eq!(graph.nodes[0].id, "11".repeat(16)); assert_eq!(graph.edges[0].source, "11".repeat(16));
787 assert_eq!(graph.edges[0].target, "22".repeat(16));
788 assert_eq!(graph.edges[0].event_at_unix_nanos, Some(123));
789
790 assert_eq!(dto.edges.len(), 2);
796 assert_eq!(dto.edges[0].kind, "caused");
797 assert_eq!(dto.edges[0].direction, "outgoing");
798 assert_eq!(dto.edges[0].peer, "42");
799 assert_eq!(dto.edges[0].weight, 0.75);
800 assert_eq!(dto.edges[0].created_at_unix_nanos, 900);
801 assert_eq!(
802 dto.edges[0].annotation.as_deref(),
803 Some("manually asserted")
804 );
805 assert_eq!(dto.edges[1].kind, "supports");
806 assert_eq!(dto.edges[1].direction, "incoming");
807 assert_eq!(dto.edges[1].peer, "7");
808 assert!(dto.edges[1].annotation.is_none());
809 }
810}