1use brain_db_sdk::wire::types::{
8 EntityCreateRequest, EntityCreateResponse, EntityGetRequest, EntityGetResponse, EntityListItem,
9 EntityListRequest, EntityResolveRequest, EntityResolveResponse, EntityView,
10 RelationTraverseRequest, RelationTraverseResponseFrame, ResolutionOutcomeWire,
11 TraversalPathWire,
12};
13use serde::{Deserialize, Serialize};
14
15use super::{parse_uuid, uuid_string};
16
17const DEFAULT_LIST_LIMIT: u32 = 100;
19const MAX_LIST_LIMIT: u32 = 1000;
21const DEFAULT_TRAVERSE_DEPTH: u32 = 3;
23const MAX_TRAVERSE_DEPTH: u32 = 5;
25const DEFAULT_TRAVERSE_NODES: u32 = 100;
27const MAX_TRAVERSE_NODES: u32 = 1000;
29
30fn merged_into_string(id: &[u8; 16]) -> Option<String> {
33 (id != &[0u8; 16]).then(|| uuid_string(id))
34}
35
36#[derive(Debug, Deserialize)]
40pub struct ResolveEntityBody {
41 pub candidate_name: String,
43 #[serde(default)]
45 pub resolution_context: String,
46 #[serde(default)]
48 pub type_hint: u32,
49 #[serde(default)]
51 pub allow_create: bool,
52}
53
54impl ResolveEntityBody {
55 pub fn to_request(&self) -> Result<EntityResolveRequest, String> {
57 let candidate_name = self.candidate_name.trim().to_owned();
58 if candidate_name.is_empty() {
59 return Err("candidate_name must not be blank".into());
60 }
61 Ok(EntityResolveRequest {
62 candidate_name,
63 resolution_context: self.resolution_context.clone(),
64 entity_type_hint: self.type_hint,
65 allow_create: self.allow_create,
66 request_id: brain_db_sdk::new_id(),
67 act_as: None,
68 })
69 }
70}
71
72#[derive(Debug, Serialize)]
74pub struct ResolveEntityResponseDto {
75 pub outcome: &'static str,
77 pub tier: u8,
79 pub confidence: f32,
81 pub entity_id: Option<String>,
84 pub candidate_ids: Vec<String>,
86}
87
88impl From<EntityResolveResponse> for ResolveEntityResponseDto {
89 fn from(r: EntityResolveResponse) -> Self {
90 let (outcome, bound) = match r.outcome {
93 ResolutionOutcomeWire::Resolved => ("resolved", true),
94 ResolutionOutcomeWire::Created => ("created", true),
95 ResolutionOutcomeWire::Ambiguous => ("ambiguous", false),
96 ResolutionOutcomeWire::NotFound => ("not_found", false),
97 };
98 Self {
99 outcome,
100 tier: r.tier,
101 confidence: r.confidence,
102 entity_id: bound.then(|| uuid_string(&r.resolved_entity)),
103 candidate_ids: r.candidate_ids.iter().map(uuid_string).collect(),
104 }
105 }
106}
107
108#[derive(Debug, Deserialize)]
112pub struct CreateEntityBody {
113 pub entity_type_id: u32,
116 pub canonical_name: String,
118 #[serde(default)]
120 pub aliases: Vec<String>,
121}
122
123impl CreateEntityBody {
124 pub fn to_request(&self) -> Result<EntityCreateRequest, String> {
126 if self.entity_type_id == 0 {
127 return Err("entity_type_id must be non-zero".into());
128 }
129 let canonical_name = self.canonical_name.trim().to_owned();
130 if canonical_name.is_empty() {
131 return Err("canonical_name must not be blank".into());
132 }
133 Ok(EntityCreateRequest {
134 entity_type_id: self.entity_type_id,
135 canonical_name,
136 aliases: self.aliases.clone(),
137 attributes_blob: Vec::new(),
138 session_id: 0,
139 request_id: brain_db_sdk::new_id(),
140 act_as: None,
141 })
142 }
143}
144
145#[derive(Debug, Serialize)]
147pub struct CreateEntityResponseDto {
148 pub entity_id: String,
150}
151
152impl From<EntityCreateResponse> for CreateEntityResponseDto {
153 fn from(r: EntityCreateResponse) -> Self {
154 Self {
155 entity_id: uuid_string(&r.entity_id),
156 }
157 }
158}
159
160#[derive(Debug, Serialize)]
164pub struct EntityDetailDto {
165 pub entity_id: String,
167 pub entity_type_id: u32,
169 pub canonical_name: String,
171 pub aliases: Vec<String>,
173 pub mention_count: u32,
175 pub created_at_unix_nanos: u64,
177 pub updated_at_unix_nanos: u64,
179 pub merged_into: Option<String>,
181 pub resolved_from: Vec<String>,
187}
188
189impl From<EntityView> for EntityDetailDto {
190 fn from(v: EntityView) -> Self {
191 Self {
192 entity_id: uuid_string(&v.entity_id),
193 entity_type_id: v.entity_type_id,
194 canonical_name: v.canonical_name,
195 aliases: v.aliases,
196 mention_count: v.mention_count,
197 created_at_unix_nanos: v.created_at_unix_nanos,
198 updated_at_unix_nanos: v.updated_at_unix_nanos,
199 merged_into: merged_into_string(&v.merged_into),
200 resolved_from: Vec::new(),
202 }
203 }
204}
205
206impl EntityDetailDto {
207 pub fn from_get_response(resp: EntityGetResponse) -> Self {
210 let resolved_from = resp.resolved_from.iter().map(uuid_string).collect();
211 Self {
212 resolved_from,
213 ..Self::from(resp.entity)
214 }
215 }
216}
217
218pub fn get_request_from_id(id: &str) -> Result<EntityGetRequest, String> {
220 Ok(EntityGetRequest {
221 entity_id: parse_uuid(id)?,
222 act_as: None,
223 })
224}
225
226#[derive(Debug, Default, Deserialize)]
237pub struct ListEntitiesQuery {
238 #[serde(default)]
241 pub type_id: u32,
242 #[serde(default)]
244 pub prefix: String,
245 #[serde(default)]
247 pub mention_count_min: u32,
248 #[serde(default)]
250 pub include_tombstoned: bool,
251 #[serde(default)]
253 pub include_merged: bool,
254 pub limit: Option<u32>,
256}
257
258impl ListEntitiesQuery {
259 pub fn to_request(&self) -> EntityListRequest {
262 let limit = self
263 .limit
264 .unwrap_or(DEFAULT_LIST_LIMIT)
265 .clamp(1, MAX_LIST_LIMIT);
266 EntityListRequest {
267 entity_type_id: self.type_id,
268 name_prefix: self.prefix.clone(),
269 mention_count_min: self.mention_count_min,
270 include_tombstoned: self.include_tombstoned,
271 include_merged: self.include_merged,
272 limit,
273 cursor: Vec::new(),
274 act_as: None,
275 }
276 }
277}
278
279#[derive(Debug, Serialize)]
281pub struct ListEntitiesResponseDto {
282 pub entities: Vec<EntityDetailDto>,
284 pub count: usize,
286}
287
288impl ListEntitiesResponseDto {
289 pub fn from_items(items: Vec<EntityListItem>) -> Self {
291 let entities: Vec<EntityDetailDto> = items
292 .into_iter()
293 .map(|item| EntityDetailDto::from(item.entity))
294 .collect();
295 let count = entities.len();
296 Self { entities, count }
297 }
298}
299
300#[derive(Debug, Default, Deserialize)]
306pub struct TraverseBody {
307 #[serde(default)]
309 pub direction: String,
310 #[serde(default)]
312 pub relation_types: Vec<String>,
313 #[serde(default)]
315 pub max_depth: Option<u32>,
316 #[serde(default)]
318 pub max_nodes: Option<u32>,
319 #[serde(default)]
321 pub include_superseded: bool,
322}
323
324fn traverse_direction_wire(raw: &str) -> Result<u8, String> {
326 match raw.trim().to_ascii_lowercase().as_str() {
327 "" | "outgoing" | "out" => Ok(0),
328 "incoming" | "in" => Ok(1),
329 "both" => Ok(2),
330 _ => Err("direction must be one of: outgoing, incoming, both".into()),
331 }
332}
333
334impl TraverseBody {
335 pub fn to_request(&self, id: &str) -> Result<RelationTraverseRequest, String> {
338 let start_entity = parse_uuid(id)?;
339 let direction = traverse_direction_wire(&self.direction)?;
340 let max_depth = self
341 .max_depth
342 .unwrap_or(DEFAULT_TRAVERSE_DEPTH)
343 .clamp(1, MAX_TRAVERSE_DEPTH);
344 let max_nodes = self
345 .max_nodes
346 .unwrap_or(DEFAULT_TRAVERSE_NODES)
347 .clamp(1, MAX_TRAVERSE_NODES);
348 Ok(RelationTraverseRequest {
349 start_entity,
350 relation_types: self.relation_types.clone(),
351 direction,
352 max_depth,
353 max_nodes,
354 time_at_unix_nanos: 0,
355 include_superseded: self.include_superseded,
356 request_id: brain_db_sdk::new_id(),
357 act_as: None,
358 })
359 }
360}
361
362#[derive(Debug, Serialize)]
364pub struct TraversalStepDto {
365 pub relation_id: String,
367 pub from: String,
369 pub to: String,
371 pub relation_type: String,
373 pub depth: u32,
375}
376
377#[derive(Debug, Serialize)]
379pub struct TraversalPathDto {
380 pub steps: Vec<TraversalStepDto>,
382}
383
384impl From<TraversalPathWire> for TraversalPathDto {
385 fn from(p: TraversalPathWire) -> Self {
386 Self {
387 steps: p
388 .steps
389 .into_iter()
390 .map(|s| TraversalStepDto {
391 relation_id: uuid_string(&s.relation_id),
392 from: uuid_string(&s.from),
393 to: uuid_string(&s.to),
394 relation_type: s.relation_type,
395 depth: s.depth,
396 })
397 .collect(),
398 }
399 }
400}
401
402#[derive(Debug, Serialize)]
404pub struct TraverseResponseDto {
405 pub paths: Vec<TraversalPathDto>,
407 pub total_paths: u32,
409 pub truncated: bool,
411}
412
413impl TraverseResponseDto {
414 pub fn from_frames(frames: Vec<RelationTraverseResponseFrame>) -> Self {
419 let truncated = frames.iter().any(|f| f.truncated);
420 let total_paths = frames.last().map_or(0, |f| f.total_paths);
421 let paths = frames
422 .into_iter()
423 .flat_map(|f| f.paths.into_iter().map(TraversalPathDto::from))
424 .collect();
425 Self {
426 paths,
427 total_paths,
428 truncated,
429 }
430 }
431}
432
433#[cfg(test)]
434mod entity_tests {
435 use super::*;
436 use brain_db_sdk::wire::types::{TraversalPathWire, TraversalStepWire};
437
438 #[test]
439 fn create_body_rejects_zero_type_and_blank_name() {
440 let bad_type = CreateEntityBody {
441 entity_type_id: 0,
442 canonical_name: "Ada".into(),
443 aliases: vec![],
444 };
445 assert!(bad_type.to_request().is_err());
446 let blank = CreateEntityBody {
447 entity_type_id: 1,
448 canonical_name: " ".into(),
449 aliases: vec![],
450 };
451 assert!(blank.to_request().is_err());
452 let ok = CreateEntityBody {
453 entity_type_id: 1,
454 canonical_name: " Ada Lovelace ".into(),
455 aliases: vec!["Ada".into()],
456 };
457 let req = ok.to_request().unwrap();
458 assert_eq!(req.canonical_name, "Ada Lovelace"); assert_eq!(req.entity_type_id, 1);
460 assert_eq!(req.attributes_blob, [] as [u8; 0]);
461 }
462
463 #[test]
464 fn entity_get_response_surfaces_resolved_from_chain() {
465 let mut survivor = [0u8; 16];
466 survivor[15] = 3; let mut a = [0u8; 16];
468 a[15] = 1;
469 let mut b = [0u8; 16];
470 b[15] = 2;
471 let view = EntityView {
472 entity_id: survivor,
473 entity_type_id: 1,
474 canonical_name: "Ada Lovelace".into(),
475 normalized_name: "ada lovelace".into(),
476 aliases: vec![],
477 attributes_blob: vec![],
478 mention_count: 4,
479 created_at_unix_nanos: 1,
480 updated_at_unix_nanos: 2,
481 merged_into: [0u8; 16],
482 embedding_version: 1,
483 flags: 0,
484 };
485 let resp = EntityGetResponse {
486 entity: view.clone(),
487 resolved_from: vec![a, b],
488 };
489 let dto = EntityDetailDto::from_get_response(resp);
490 assert_eq!(dto.entity_id, uuid_string(&survivor));
491 assert_eq!(dto.resolved_from, vec![uuid_string(&a), uuid_string(&b)]);
493 assert_eq!(EntityDetailDto::from(view).resolved_from, [] as [String; 0]);
495 }
496
497 #[test]
498 fn resolve_body_maps_type_hint_and_rejects_blank() {
499 assert!(
500 ResolveEntityBody {
501 candidate_name: " ".into(),
502 resolution_context: String::new(),
503 type_hint: 0,
504 allow_create: false,
505 }
506 .to_request()
507 .is_err()
508 );
509 let req = ResolveEntityBody {
510 candidate_name: "Ada".into(),
511 resolution_context: "the mathematician".into(),
512 type_hint: 5,
513 allow_create: true,
514 }
515 .to_request()
516 .unwrap();
517 assert_eq!(req.entity_type_hint, 5);
518 assert!(req.allow_create);
519 assert_eq!(req.resolution_context, "the mathematician");
520 }
521
522 #[test]
523 fn resolve_response_surfaces_id_only_when_bound() {
524 let mut id = [0u8; 16];
525 id[15] = 9;
526 let resolved = EntityResolveResponse {
527 outcome: ResolutionOutcomeWire::Resolved,
528 tier: 1,
529 confidence: 1.0,
530 resolved_entity: id,
531 candidate_ids: vec![],
532 audit_id: [0u8; 16],
533 };
534 let dto = ResolveEntityResponseDto::from(resolved);
535 assert_eq!(dto.outcome, "resolved");
536 assert_eq!(
537 dto.entity_id.as_deref(),
538 Some("00000000-0000-0000-0000-000000000009")
539 );
540
541 let ambiguous = EntityResolveResponse {
542 outcome: ResolutionOutcomeWire::Ambiguous,
543 tier: 0,
544 confidence: 0.0,
545 resolved_entity: id, candidate_ids: vec![[0x11u8; 16], [0x22u8; 16]],
547 audit_id: [0u8; 16],
548 };
549 let dto = ResolveEntityResponseDto::from(ambiguous);
550 assert_eq!(dto.outcome, "ambiguous");
551 assert!(dto.entity_id.is_none());
552 assert_eq!(dto.candidate_ids.len(), 2);
553 }
554
555 #[test]
556 fn list_query_clamps_limit_and_defaults() {
557 let q = ListEntitiesQuery {
558 limit: Some(9999),
559 ..Default::default()
560 };
561 let req = q.to_request();
562 assert_eq!(req.limit, MAX_LIST_LIMIT);
563 assert_eq!(req.cursor, [] as [u8; 0]);
564 let none = ListEntitiesQuery::default().to_request();
565 assert_eq!(none.limit, DEFAULT_LIST_LIMIT);
566 }
567
568 #[test]
569 fn traverse_body_parses_direction_and_clamps() {
570 let id = "00000000-0000-0000-0000-000000000001";
571 let req = TraverseBody {
572 direction: "both".into(),
573 relation_types: vec!["brain:works_at".into()],
574 max_depth: Some(99),
575 max_nodes: Some(0),
576 include_superseded: true,
577 }
578 .to_request(id)
579 .unwrap();
580 assert_eq!(req.direction, 2);
581 assert_eq!(req.max_depth, MAX_TRAVERSE_DEPTH);
582 assert_eq!(req.max_nodes, 1); assert!(req.include_superseded);
584
585 let req = TraverseBody::default().to_request(id).unwrap();
587 assert_eq!(req.direction, 0);
588 assert_eq!(req.max_depth, DEFAULT_TRAVERSE_DEPTH);
589
590 assert!(
591 TraverseBody {
592 direction: "sideways".into(),
593 ..Default::default()
594 }
595 .to_request(id)
596 .is_err()
597 );
598 assert!(TraverseBody::default().to_request("not-a-uuid").is_err());
599 }
600
601 #[test]
602 fn traverse_response_folds_frames_and_uuidifies() {
603 let step = TraversalStepWire {
604 relation_id: [0x01u8; 16],
605 from: [0x02u8; 16],
606 to: [0x03u8; 16],
607 relation_type: "brain:works_at".into(),
608 depth: 1,
609 };
610 let frames = vec![
611 RelationTraverseResponseFrame {
612 paths: vec![TraversalPathWire {
613 steps: vec![step.clone()],
614 }],
615 total_paths: 2,
616 truncated: false,
617 is_final: false,
618 },
619 RelationTraverseResponseFrame {
620 paths: vec![TraversalPathWire { steps: vec![step] }],
621 total_paths: 2,
622 truncated: true,
623 is_final: true,
624 },
625 ];
626 let dto = TraverseResponseDto::from_frames(frames);
627 assert_eq!(dto.paths.len(), 2);
628 assert_eq!(dto.total_paths, 2);
629 assert!(dto.truncated);
630 assert!(dto.paths[0].steps[0].from.contains('-'));
631 }
632}