#![allow(missing_docs)]
use brain_db_sdk::wire::types::{
EdgeKindWire, GraphEdge, GraphFetchRequest, GraphFetchResponseFrame, GraphNode, LinkResponse,
UnlinkResponse, WireMemoryId,
};
use serde::{Deserialize, Serialize};
use super::{hex_decode, hex_encode, mem_id_decimal, parse_memory_id, uuid_string};
fn parse_edge_kind(s: &str) -> Result<EdgeKindWire, String> {
Ok(match s.trim().to_ascii_lowercase().as_str() {
"caused" => EdgeKindWire::Caused,
"followed_by" | "followedby" => EdgeKindWire::FollowedBy,
"derived_from" | "derivedfrom" => EdgeKindWire::DerivedFrom,
"similar_to" | "similarto" => EdgeKindWire::SimilarTo,
"contradicts" => EdgeKindWire::Contradicts,
"supports" => EdgeKindWire::Supports,
"references" => EdgeKindWire::References,
"part_of" | "partof" => EdgeKindWire::PartOf,
other => return Err(format!("unknown edge kind `{other}`")),
})
}
pub(crate) fn edge_kind_str(k: EdgeKindWire) -> &'static str {
match k {
EdgeKindWire::Caused => "caused",
EdgeKindWire::FollowedBy => "followed_by",
EdgeKindWire::DerivedFrom => "derived_from",
EdgeKindWire::SimilarTo => "similar_to",
EdgeKindWire::Contradicts => "contradicts",
EdgeKindWire::Supports => "supports",
EdgeKindWire::References => "references",
EdgeKindWire::PartOf => "part_of",
}
}
#[derive(Debug, Deserialize)]
pub struct LinkBody {
pub source: String,
pub target: String,
pub kind: String,
#[serde(default = "default_weight")]
pub weight: f32,
}
fn default_weight() -> f32 {
1.0
}
impl LinkBody {
pub fn parts(&self) -> Result<(WireMemoryId, WireMemoryId, EdgeKindWire, f32), String> {
Ok((
parse_memory_id(&self.source)?,
parse_memory_id(&self.target)?,
parse_edge_kind(&self.kind)?,
self.weight,
))
}
}
#[derive(Debug, Serialize)]
pub struct LinkResponseDto {
pub source: String,
pub target: String,
pub kind: &'static str,
pub weight: f32,
pub created_at_unix_nanos: u64,
pub already_existed: bool,
}
impl From<LinkResponse> for LinkResponseDto {
fn from(r: LinkResponse) -> Self {
Self {
source: r.source.to_string(),
target: r.target.to_string(),
kind: edge_kind_str(r.kind),
weight: r.weight,
created_at_unix_nanos: r.created_at_unix_nanos,
already_existed: r.already_existed,
}
}
}
#[derive(Debug, Deserialize)]
pub struct UnlinkBody {
pub source: String,
pub target: String,
pub kind: String,
}
impl UnlinkBody {
pub fn parts(&self) -> Result<(WireMemoryId, WireMemoryId, EdgeKindWire), String> {
Ok((
parse_memory_id(&self.source)?,
parse_memory_id(&self.target)?,
parse_edge_kind(&self.kind)?,
))
}
}
#[derive(Debug, Serialize)]
pub struct UnlinkResponseDto {
pub source: String,
pub target: String,
pub kind: &'static str,
pub removed: bool,
}
impl From<UnlinkResponse> for UnlinkResponseDto {
fn from(r: UnlinkResponse) -> Self {
Self {
source: r.source.to_string(),
target: r.target.to_string(),
kind: edge_kind_str(r.kind),
removed: r.removed,
}
}
}
const GRAPH_DEFAULT_LIMIT: u32 = 200;
const GRAPH_MAX_LIMIT: u32 = 500;
const NODE_ENTITY: u8 = 0;
const NODE_STATEMENT: u8 = 1;
const NODE_MEMORY: u8 = 2;
const EDGE_MENTIONS: u8 = 3;
const EDGE_MEMORY_BASE: u8 = 4;
const EDGE_MEMORY_LAST: u8 = EDGE_MEMORY_BASE + 7;
fn graph_edge_kind_str(kind: u8) -> Option<&'static str> {
Some(match kind {
0 => "relation",
1 => "fact",
2 => "has_statement",
EDGE_MENTIONS => "mentions",
EDGE_MEMORY_BASE..=EDGE_MEMORY_LAST => {
return memory_edge_kind(kind).map(edge_kind_str);
}
_ => return None,
})
}
fn memory_edge_kind(kind: u8) -> Option<EdgeKindWire> {
if !(EDGE_MEMORY_BASE..=EDGE_MEMORY_LAST).contains(&kind) {
return None;
}
Some(match kind - EDGE_MEMORY_BASE {
0 => EdgeKindWire::Caused,
1 => EdgeKindWire::FollowedBy,
2 => EdgeKindWire::DerivedFrom,
3 => EdgeKindWire::SimilarTo,
4 => EdgeKindWire::Contradicts,
5 => EdgeKindWire::Supports,
6 => EdgeKindWire::References,
_ => EdgeKindWire::PartOf,
})
}
fn node_id_str(id: &[u8; 16], kind: u8) -> String {
if kind == NODE_MEMORY {
mem_id_decimal(id)
} else {
uuid_string(id)
}
}
#[derive(Debug, Deserialize)]
pub struct GraphFetchQuery {
#[serde(default)]
pub limit: Option<u32>,
#[serde(default)]
pub cursor: Option<String>,
#[serde(default)]
pub include_statements: Option<bool>,
#[serde(default)]
pub include_memories: Option<bool>,
#[serde(default)]
pub include_memory_edges: Option<bool>,
#[serde(default)]
pub include_tombstoned: Option<bool>,
}
impl GraphFetchQuery {
pub fn to_request(&self) -> Result<GraphFetchRequest, String> {
let limit = self
.limit
.unwrap_or(GRAPH_DEFAULT_LIMIT)
.clamp(1, GRAPH_MAX_LIMIT);
let cursor = match &self.cursor {
Some(c) if !c.is_empty() => hex_decode(c)?,
_ => Vec::new(),
};
let include_memories = self.include_memories.unwrap_or(false);
let include_memory_edges = self.include_memory_edges.unwrap_or(false);
if include_memory_edges && !include_memories {
return Err(
"include_memory_edges requires include_memories: memory↔memory edges hang off \
memory nodes, and without that layer their endpoints would not be in the page"
.to_string(),
);
}
Ok(GraphFetchRequest {
limit,
cursor,
include_statements: self.include_statements.unwrap_or(false),
include_memories,
include_memory_edges,
include_tombstoned: self.include_tombstoned.unwrap_or(false),
act_as: None,
})
}
}
#[derive(Debug, Serialize)]
pub struct GraphNodeDto {
pub id: String,
pub kind: &'static str,
pub label: String,
pub type_qname: String,
}
impl From<GraphNode> for GraphNodeDto {
fn from(n: GraphNode) -> Self {
let kind = match n.kind {
NODE_ENTITY => "entity",
NODE_STATEMENT => "statement",
NODE_MEMORY => "memory",
_ => "unknown",
};
Self {
id: node_id_str(&n.id, n.kind),
kind,
label: n.label,
type_qname: n.type_qname,
}
}
}
#[derive(Debug, Serialize)]
pub struct GraphEdgeDto {
pub from_id: String,
pub to_id: String,
pub kind: &'static str,
pub label: String,
}
impl From<GraphEdge> for GraphEdgeDto {
fn from(e: GraphEdge) -> Self {
let kind = graph_edge_kind_str(e.kind).unwrap_or("unknown");
let (from_kind, to_kind) = if e.kind == EDGE_MENTIONS {
(NODE_MEMORY, NODE_ENTITY)
} else if memory_edge_kind(e.kind).is_some() {
(NODE_MEMORY, NODE_MEMORY)
} else {
(NODE_ENTITY, NODE_ENTITY)
};
Self {
from_id: node_id_str(&e.from_id, from_kind),
to_id: node_id_str(&e.to_id, to_kind),
kind,
label: e.label,
}
}
}
#[derive(Debug, Serialize)]
pub struct GraphPageDto {
pub nodes: Vec<GraphNodeDto>,
pub edges: Vec<GraphEdgeDto>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
}
impl GraphPageDto {
pub fn from_frames(frames: Vec<GraphFetchResponseFrame>) -> Self {
let next = frames
.last()
.map(|f| f.next_cursor.clone())
.filter(|c| !c.is_empty())
.map(|c| hex_encode(&c));
let mut nodes = Vec::new();
let mut edges = Vec::new();
for f in frames {
nodes.extend(f.nodes.into_iter().map(GraphNodeDto::from));
edges.extend(f.edges.into_iter().map(GraphEdgeDto::from));
}
Self {
nodes,
edges,
next_cursor: next,
}
}
}
#[cfg(test)]
mod graph_tests {
use super::*;
#[test]
fn query_clamps_limit_and_defaults_layers() {
let q = GraphFetchQuery {
limit: Some(9999),
cursor: None,
include_statements: None,
include_memories: Some(true),
include_memory_edges: None,
include_tombstoned: None,
};
let r = q.to_request().unwrap();
assert_eq!(r.limit, GRAPH_MAX_LIMIT);
assert!(!r.include_statements);
assert!(r.include_memories);
assert!(!r.include_memory_edges);
assert_eq!(r.cursor, [] as [u8; 0]);
}
#[test]
fn memory_edges_require_the_memory_layer() {
let alone = GraphFetchQuery {
limit: None,
cursor: None,
include_statements: None,
include_memories: None,
include_memory_edges: Some(true),
include_tombstoned: None,
};
let err = alone.to_request().unwrap_err();
assert!(err.contains("include_memories"), "{err}");
let together = GraphFetchQuery {
limit: None,
cursor: None,
include_statements: None,
include_memories: Some(true),
include_memory_edges: Some(true),
include_tombstoned: None,
};
let r = together.to_request().unwrap();
assert!(r.include_memories && r.include_memory_edges);
}
#[test]
fn memory_edge_bytes_name_each_kind_and_read_both_ends_as_memories() {
let mut a = [0u8; 16];
a[15] = 11;
let mut b = [0u8; 16];
b[15] = 12;
let expected = [
(4u8, "caused"),
(5, "followed_by"),
(6, "derived_from"),
(7, "similar_to"),
(8, "contradicts"),
(9, "supports"),
(10, "references"),
(11, "part_of"),
];
for (byte, name) in expected {
let e = GraphEdgeDto::from(GraphEdge {
from_id: a,
to_id: b,
kind: byte,
label: name.to_string(),
});
assert_eq!(e.kind, name, "byte {byte}");
assert_eq!(e.from_id, "11");
assert_eq!(e.to_id, "12");
}
}
#[test]
fn unknown_edge_byte_is_not_guessed() {
let e = GraphEdgeDto::from(GraphEdge {
from_id: [0u8; 16],
to_id: [0u8; 16],
kind: 99,
label: String::new(),
});
assert_eq!(e.kind, "unknown");
}
#[test]
fn entity_node_is_uuid_memory_node_is_decimal() {
let mut id = [0u8; 16];
id[15] = 5;
let ent = GraphNodeDto::from(GraphNode {
id,
kind: NODE_ENTITY,
label: "Sarah Chen".into(),
type_qname: "brain:Person".into(),
});
assert_eq!(ent.kind, "entity");
assert_eq!(ent.id, "00000000-0000-0000-0000-000000000005");
let mem = GraphNodeDto::from(GraphNode {
id,
kind: NODE_MEMORY,
label: "a memory".into(),
type_qname: String::new(),
});
assert_eq!(mem.kind, "memory");
assert_eq!(mem.id, "5"); }
#[test]
fn mentions_edge_source_is_decimal_memory() {
let mut mem = [0u8; 16];
mem[15] = 7;
let ent = [0x11u8; 16];
let e = GraphEdgeDto::from(GraphEdge {
from_id: mem,
to_id: ent,
kind: EDGE_MENTIONS,
label: String::new(),
});
assert_eq!(e.kind, "mentions");
assert_eq!(e.from_id, "7"); assert!(e.to_id.contains('-')); }
#[test]
fn relation_edge_both_ends_uuid() {
let a = [0x01u8; 16];
let b = [0x02u8; 16];
let e = GraphEdgeDto::from(GraphEdge {
from_id: a,
to_id: b,
kind: 0,
label: "brain:works_at".into(),
});
assert_eq!(e.kind, "relation");
assert!(e.from_id.contains('-') && e.to_id.contains('-'));
}
}