use brain_db_sdk::wire::types::{
EntityCreateRequest, EntityCreateResponse, EntityGetRequest, EntityGetResponse, EntityListItem,
EntityListRequest, EntityResolveRequest, EntityResolveResponse, EntityView,
RelationTraverseRequest, RelationTraverseResponseFrame, ResolutionOutcomeWire,
TraversalPathWire,
};
use serde::{Deserialize, Serialize};
use super::{parse_uuid, uuid_string};
const DEFAULT_LIST_LIMIT: u32 = 100;
const MAX_LIST_LIMIT: u32 = 1000;
const DEFAULT_TRAVERSE_DEPTH: u32 = 3;
const MAX_TRAVERSE_DEPTH: u32 = 5;
const DEFAULT_TRAVERSE_NODES: u32 = 100;
const MAX_TRAVERSE_NODES: u32 = 1000;
fn merged_into_string(id: &[u8; 16]) -> Option<String> {
(id != &[0u8; 16]).then(|| uuid_string(id))
}
#[derive(Debug, Deserialize)]
pub struct ResolveEntityBody {
pub candidate_name: String,
#[serde(default)]
pub resolution_context: String,
#[serde(default)]
pub type_hint: u32,
#[serde(default)]
pub allow_create: bool,
}
impl ResolveEntityBody {
pub fn to_request(&self) -> Result<EntityResolveRequest, String> {
let candidate_name = self.candidate_name.trim().to_owned();
if candidate_name.is_empty() {
return Err("candidate_name must not be blank".into());
}
Ok(EntityResolveRequest {
candidate_name,
resolution_context: self.resolution_context.clone(),
entity_type_hint: self.type_hint,
allow_create: self.allow_create,
request_id: brain_db_sdk::new_id(),
act_as: None,
})
}
}
#[derive(Debug, Serialize)]
pub struct ResolveEntityResponseDto {
pub outcome: &'static str,
pub tier: u8,
pub confidence: f32,
pub entity_id: Option<String>,
pub candidate_ids: Vec<String>,
}
impl From<EntityResolveResponse> for ResolveEntityResponseDto {
fn from(r: EntityResolveResponse) -> Self {
let (outcome, bound) = match r.outcome {
ResolutionOutcomeWire::Resolved => ("resolved", true),
ResolutionOutcomeWire::Created => ("created", true),
ResolutionOutcomeWire::Ambiguous => ("ambiguous", false),
ResolutionOutcomeWire::NotFound => ("not_found", false),
};
Self {
outcome,
tier: r.tier,
confidence: r.confidence,
entity_id: bound.then(|| uuid_string(&r.resolved_entity)),
candidate_ids: r.candidate_ids.iter().map(uuid_string).collect(),
}
}
}
#[derive(Debug, Deserialize)]
pub struct CreateEntityBody {
pub entity_type_id: u32,
pub canonical_name: String,
#[serde(default)]
pub aliases: Vec<String>,
}
impl CreateEntityBody {
pub fn to_request(&self) -> Result<EntityCreateRequest, String> {
if self.entity_type_id == 0 {
return Err("entity_type_id must be non-zero".into());
}
let canonical_name = self.canonical_name.trim().to_owned();
if canonical_name.is_empty() {
return Err("canonical_name must not be blank".into());
}
Ok(EntityCreateRequest {
entity_type_id: self.entity_type_id,
canonical_name,
aliases: self.aliases.clone(),
attributes_blob: Vec::new(),
session_id: 0,
request_id: brain_db_sdk::new_id(),
act_as: None,
})
}
}
#[derive(Debug, Serialize)]
pub struct CreateEntityResponseDto {
pub entity_id: String,
}
impl From<EntityCreateResponse> for CreateEntityResponseDto {
fn from(r: EntityCreateResponse) -> Self {
Self {
entity_id: uuid_string(&r.entity_id),
}
}
}
#[derive(Debug, Serialize)]
pub struct EntityDetailDto {
pub entity_id: String,
pub entity_type_id: u32,
pub canonical_name: String,
pub aliases: Vec<String>,
pub mention_count: u32,
pub created_at_unix_nanos: u64,
pub updated_at_unix_nanos: u64,
pub merged_into: Option<String>,
pub resolved_from: Vec<String>,
}
impl From<EntityView> for EntityDetailDto {
fn from(v: EntityView) -> Self {
Self {
entity_id: uuid_string(&v.entity_id),
entity_type_id: v.entity_type_id,
canonical_name: v.canonical_name,
aliases: v.aliases,
mention_count: v.mention_count,
created_at_unix_nanos: v.created_at_unix_nanos,
updated_at_unix_nanos: v.updated_at_unix_nanos,
merged_into: merged_into_string(&v.merged_into),
resolved_from: Vec::new(),
}
}
}
impl EntityDetailDto {
pub fn from_get_response(resp: EntityGetResponse) -> Self {
let resolved_from = resp.resolved_from.iter().map(uuid_string).collect();
Self {
resolved_from,
..Self::from(resp.entity)
}
}
}
pub fn get_request_from_id(id: &str) -> Result<EntityGetRequest, String> {
Ok(EntityGetRequest {
entity_id: parse_uuid(id)?,
act_as: None,
})
}
#[derive(Debug, Default, Deserialize)]
pub struct ListEntitiesQuery {
#[serde(default)]
pub type_id: u32,
#[serde(default)]
pub prefix: String,
#[serde(default)]
pub mention_count_min: u32,
#[serde(default)]
pub include_tombstoned: bool,
#[serde(default)]
pub include_merged: bool,
pub limit: Option<u32>,
}
impl ListEntitiesQuery {
pub fn to_request(&self) -> EntityListRequest {
let limit = self
.limit
.unwrap_or(DEFAULT_LIST_LIMIT)
.clamp(1, MAX_LIST_LIMIT);
EntityListRequest {
entity_type_id: self.type_id,
name_prefix: self.prefix.clone(),
mention_count_min: self.mention_count_min,
include_tombstoned: self.include_tombstoned,
include_merged: self.include_merged,
limit,
cursor: Vec::new(),
act_as: None,
}
}
}
#[derive(Debug, Serialize)]
pub struct ListEntitiesResponseDto {
pub entities: Vec<EntityDetailDto>,
pub count: usize,
}
impl ListEntitiesResponseDto {
pub fn from_items(items: Vec<EntityListItem>) -> Self {
let entities: Vec<EntityDetailDto> = items
.into_iter()
.map(|item| EntityDetailDto::from(item.entity))
.collect();
let count = entities.len();
Self { entities, count }
}
}
#[derive(Debug, Default, Deserialize)]
pub struct TraverseBody {
#[serde(default)]
pub direction: String,
#[serde(default)]
pub relation_types: Vec<String>,
#[serde(default)]
pub max_depth: Option<u32>,
#[serde(default)]
pub max_nodes: Option<u32>,
#[serde(default)]
pub include_superseded: bool,
}
fn traverse_direction_wire(raw: &str) -> Result<u8, String> {
match raw.trim().to_ascii_lowercase().as_str() {
"" | "outgoing" | "out" => Ok(0),
"incoming" | "in" => Ok(1),
"both" => Ok(2),
_ => Err("direction must be one of: outgoing, incoming, both".into()),
}
}
impl TraverseBody {
pub fn to_request(&self, id: &str) -> Result<RelationTraverseRequest, String> {
let start_entity = parse_uuid(id)?;
let direction = traverse_direction_wire(&self.direction)?;
let max_depth = self
.max_depth
.unwrap_or(DEFAULT_TRAVERSE_DEPTH)
.clamp(1, MAX_TRAVERSE_DEPTH);
let max_nodes = self
.max_nodes
.unwrap_or(DEFAULT_TRAVERSE_NODES)
.clamp(1, MAX_TRAVERSE_NODES);
Ok(RelationTraverseRequest {
start_entity,
relation_types: self.relation_types.clone(),
direction,
max_depth,
max_nodes,
time_at_unix_nanos: 0,
include_superseded: self.include_superseded,
request_id: brain_db_sdk::new_id(),
act_as: None,
})
}
}
#[derive(Debug, Serialize)]
pub struct TraversalStepDto {
pub relation_id: String,
pub from: String,
pub to: String,
pub relation_type: String,
pub depth: u32,
}
#[derive(Debug, Serialize)]
pub struct TraversalPathDto {
pub steps: Vec<TraversalStepDto>,
}
impl From<TraversalPathWire> for TraversalPathDto {
fn from(p: TraversalPathWire) -> Self {
Self {
steps: p
.steps
.into_iter()
.map(|s| TraversalStepDto {
relation_id: uuid_string(&s.relation_id),
from: uuid_string(&s.from),
to: uuid_string(&s.to),
relation_type: s.relation_type,
depth: s.depth,
})
.collect(),
}
}
}
#[derive(Debug, Serialize)]
pub struct TraverseResponseDto {
pub paths: Vec<TraversalPathDto>,
pub total_paths: u32,
pub truncated: bool,
}
impl TraverseResponseDto {
pub fn from_frames(frames: Vec<RelationTraverseResponseFrame>) -> Self {
let truncated = frames.iter().any(|f| f.truncated);
let total_paths = frames.last().map_or(0, |f| f.total_paths);
let paths = frames
.into_iter()
.flat_map(|f| f.paths.into_iter().map(TraversalPathDto::from))
.collect();
Self {
paths,
total_paths,
truncated,
}
}
}
#[cfg(test)]
mod entity_tests {
use super::*;
use brain_db_sdk::wire::types::{TraversalPathWire, TraversalStepWire};
#[test]
fn create_body_rejects_zero_type_and_blank_name() {
let bad_type = CreateEntityBody {
entity_type_id: 0,
canonical_name: "Ada".into(),
aliases: vec![],
};
assert!(bad_type.to_request().is_err());
let blank = CreateEntityBody {
entity_type_id: 1,
canonical_name: " ".into(),
aliases: vec![],
};
assert!(blank.to_request().is_err());
let ok = CreateEntityBody {
entity_type_id: 1,
canonical_name: " Ada Lovelace ".into(),
aliases: vec!["Ada".into()],
};
let req = ok.to_request().unwrap();
assert_eq!(req.canonical_name, "Ada Lovelace"); assert_eq!(req.entity_type_id, 1);
assert_eq!(req.attributes_blob, [] as [u8; 0]);
}
#[test]
fn entity_get_response_surfaces_resolved_from_chain() {
let mut survivor = [0u8; 16];
survivor[15] = 3; let mut a = [0u8; 16];
a[15] = 1;
let mut b = [0u8; 16];
b[15] = 2;
let view = EntityView {
entity_id: survivor,
entity_type_id: 1,
canonical_name: "Ada Lovelace".into(),
normalized_name: "ada lovelace".into(),
aliases: vec![],
attributes_blob: vec![],
mention_count: 4,
created_at_unix_nanos: 1,
updated_at_unix_nanos: 2,
merged_into: [0u8; 16],
embedding_version: 1,
flags: 0,
};
let resp = EntityGetResponse {
entity: view.clone(),
resolved_from: vec![a, b],
};
let dto = EntityDetailDto::from_get_response(resp);
assert_eq!(dto.entity_id, uuid_string(&survivor));
assert_eq!(dto.resolved_from, vec![uuid_string(&a), uuid_string(&b)]);
assert_eq!(EntityDetailDto::from(view).resolved_from, [] as [String; 0]);
}
#[test]
fn resolve_body_maps_type_hint_and_rejects_blank() {
assert!(
ResolveEntityBody {
candidate_name: " ".into(),
resolution_context: String::new(),
type_hint: 0,
allow_create: false,
}
.to_request()
.is_err()
);
let req = ResolveEntityBody {
candidate_name: "Ada".into(),
resolution_context: "the mathematician".into(),
type_hint: 5,
allow_create: true,
}
.to_request()
.unwrap();
assert_eq!(req.entity_type_hint, 5);
assert!(req.allow_create);
assert_eq!(req.resolution_context, "the mathematician");
}
#[test]
fn resolve_response_surfaces_id_only_when_bound() {
let mut id = [0u8; 16];
id[15] = 9;
let resolved = EntityResolveResponse {
outcome: ResolutionOutcomeWire::Resolved,
tier: 1,
confidence: 1.0,
resolved_entity: id,
candidate_ids: vec![],
audit_id: [0u8; 16],
};
let dto = ResolveEntityResponseDto::from(resolved);
assert_eq!(dto.outcome, "resolved");
assert_eq!(
dto.entity_id.as_deref(),
Some("00000000-0000-0000-0000-000000000009")
);
let ambiguous = EntityResolveResponse {
outcome: ResolutionOutcomeWire::Ambiguous,
tier: 0,
confidence: 0.0,
resolved_entity: id, candidate_ids: vec![[0x11u8; 16], [0x22u8; 16]],
audit_id: [0u8; 16],
};
let dto = ResolveEntityResponseDto::from(ambiguous);
assert_eq!(dto.outcome, "ambiguous");
assert!(dto.entity_id.is_none());
assert_eq!(dto.candidate_ids.len(), 2);
}
#[test]
fn list_query_clamps_limit_and_defaults() {
let q = ListEntitiesQuery {
limit: Some(9999),
..Default::default()
};
let req = q.to_request();
assert_eq!(req.limit, MAX_LIST_LIMIT);
assert_eq!(req.cursor, [] as [u8; 0]);
let none = ListEntitiesQuery::default().to_request();
assert_eq!(none.limit, DEFAULT_LIST_LIMIT);
}
#[test]
fn traverse_body_parses_direction_and_clamps() {
let id = "00000000-0000-0000-0000-000000000001";
let req = TraverseBody {
direction: "both".into(),
relation_types: vec!["brain:works_at".into()],
max_depth: Some(99),
max_nodes: Some(0),
include_superseded: true,
}
.to_request(id)
.unwrap();
assert_eq!(req.direction, 2);
assert_eq!(req.max_depth, MAX_TRAVERSE_DEPTH);
assert_eq!(req.max_nodes, 1); assert!(req.include_superseded);
let req = TraverseBody::default().to_request(id).unwrap();
assert_eq!(req.direction, 0);
assert_eq!(req.max_depth, DEFAULT_TRAVERSE_DEPTH);
assert!(
TraverseBody {
direction: "sideways".into(),
..Default::default()
}
.to_request(id)
.is_err()
);
assert!(TraverseBody::default().to_request("not-a-uuid").is_err());
}
#[test]
fn traverse_response_folds_frames_and_uuidifies() {
let step = TraversalStepWire {
relation_id: [0x01u8; 16],
from: [0x02u8; 16],
to: [0x03u8; 16],
relation_type: "brain:works_at".into(),
depth: 1,
};
let frames = vec![
RelationTraverseResponseFrame {
paths: vec![TraversalPathWire {
steps: vec![step.clone()],
}],
total_paths: 2,
truncated: false,
is_final: false,
},
RelationTraverseResponseFrame {
paths: vec![TraversalPathWire { steps: vec![step] }],
total_paths: 2,
truncated: true,
is_final: true,
},
];
let dto = TraverseResponseDto::from_frames(frames);
assert_eq!(dto.paths.len(), 2);
assert_eq!(dto.total_paths, 2);
assert!(dto.truncated);
assert!(dto.paths[0].steps[0].from.contains('-'));
}
}