Skip to main content

ipfrs_semantic/
knowledge_base_builder.rs

1//! # Knowledge Base Builder
2//!
3//! Incrementally builds and maintains a semantic knowledge base from documents,
4//! supporting semantic triples, entity relationships, and concept graphs.
5//!
6//! ## Features
7//!
8//! - Entity CRUD with alias indexing and optional embedding vectors
9//! - Relation triples with FNV-1a-based IDs and confidence scores
10//! - Concept co-occurrence graphs derived from document ingestion
11//! - BFS path-finding between entities over the relation graph
12//! - Rich statistics for monitoring and inspection
13
14use std::collections::{HashMap, HashSet, VecDeque};
15
16/// Error types for `KnowledgeBaseBuilder` operations.
17#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
18pub enum KbError {
19    /// An entity with the given ID already exists.
20    #[error("entity already exists: {0}")]
21    EntityAlreadyExists(String),
22    /// No entity found for the given ID.
23    #[error("entity not found: {0}")]
24    EntityNotFound(String),
25    /// A document with the given ID already exists.
26    #[error("document already exists: {0}")]
27    DocumentAlreadyExists(String),
28    /// A relation with the given ID already exists.
29    #[error("relation already exists: {0}")]
30    RelationAlreadyExists(String),
31}
32
33// ─────────────────────────────────────────────────────────────────────────────
34// Core data types
35// ─────────────────────────────────────────────────────────────────────────────
36
37/// A named entity in the knowledge base.
38#[derive(Debug, Clone)]
39pub struct KbBuilderEntity {
40    /// Unique string identifier.
41    pub id: String,
42    /// Human-readable primary name.
43    pub name: String,
44    /// Semantic type, e.g. "person", "organization", "concept".
45    pub entity_type: String,
46    /// Alternative surface forms for this entity.
47    pub aliases: Vec<String>,
48    /// Optional embedding vector.
49    pub embedding: Option<Vec<f64>>,
50    /// Unix epoch seconds at creation time.
51    pub created_at: u64,
52    /// Unix epoch seconds at last update time.
53    pub updated_at: u64,
54}
55
56impl KbBuilderEntity {
57    /// Create a minimal entity with the required fields.
58    pub fn new(
59        id: impl Into<String>,
60        name: impl Into<String>,
61        entity_type: impl Into<String>,
62        now: u64,
63    ) -> Self {
64        Self {
65            id: id.into(),
66            name: name.into(),
67            entity_type: entity_type.into(),
68            aliases: Vec::new(),
69            embedding: None,
70            created_at: now,
71            updated_at: now,
72        }
73    }
74}
75
76/// A directed semantic relation between two entities.
77#[derive(Debug, Clone)]
78pub struct KbRelation {
79    /// FNV-1a hash of `subject_id + predicate + object_id`.
80    pub id: String,
81    /// Subject entity ID.
82    pub subject_id: String,
83    /// Predicate / relationship type.
84    pub predicate: String,
85    /// Object entity ID.
86    pub object_id: String,
87    /// Confidence in [0.0, 1.0].
88    pub confidence: f64,
89    /// Provenance source identifier.
90    pub source: String,
91    /// Unix epoch seconds at creation time.
92    pub created_at: u64,
93}
94
95/// Convenience input type for a semantic triple.
96#[derive(Debug, Clone)]
97pub struct KbTriple {
98    /// Subject entity ID.
99    pub subject: String,
100    /// Predicate string.
101    pub predicate: String,
102    /// Object entity ID.
103    pub object: String,
104}
105
106impl KbTriple {
107    /// Construct a triple.
108    pub fn new(
109        subject: impl Into<String>,
110        predicate: impl Into<String>,
111        object: impl Into<String>,
112    ) -> Self {
113        Self {
114            subject: subject.into(),
115            predicate: predicate.into(),
116            object: object.into(),
117        }
118    }
119}
120
121/// A node in the concept co-occurrence graph.
122#[derive(Debug, Clone)]
123pub struct KbConceptNode {
124    /// The concept string.
125    pub concept: String,
126    /// How many documents mention this concept.
127    pub frequency: u32,
128    /// Concepts that appear in documents alongside this one.
129    pub related_concepts: Vec<String>,
130    /// IDs of documents that mention this concept.
131    pub documents: Vec<String>,
132}
133
134impl KbConceptNode {
135    fn new(concept: impl Into<String>) -> Self {
136        Self {
137            concept: concept.into(),
138            frequency: 0,
139            related_concepts: Vec::new(),
140            documents: Vec::new(),
141        }
142    }
143}
144
145/// A document ingested into the knowledge base.
146#[derive(Debug, Clone)]
147pub struct KbDocument {
148    /// Unique document identifier.
149    pub doc_id: String,
150    /// Raw text content.
151    pub content: String,
152    /// IDs of entities mentioned in this document.
153    pub entities_mentioned: Vec<String>,
154    /// Concepts extracted from this document.
155    pub concepts: Vec<String>,
156    /// Unix epoch seconds when the document was added.
157    pub added_at: u64,
158}
159
160impl KbDocument {
161    /// Construct a document record.
162    pub fn new(
163        doc_id: impl Into<String>,
164        content: impl Into<String>,
165        entities_mentioned: Vec<String>,
166        concepts: Vec<String>,
167        added_at: u64,
168    ) -> Self {
169        Self {
170            doc_id: doc_id.into(),
171            content: content.into(),
172            entities_mentioned,
173            concepts,
174            added_at,
175        }
176    }
177}
178
179/// Aggregate statistics for the knowledge base.
180#[derive(Debug, Clone)]
181pub struct KbStats {
182    /// Total entity count.
183    pub entity_count: usize,
184    /// Total relation count.
185    pub relation_count: usize,
186    /// Total document count.
187    pub document_count: usize,
188    /// Total distinct concept count.
189    pub concept_count: usize,
190    /// Average number of relations per entity (incoming + outgoing).
191    pub avg_relations_per_entity: f64,
192    /// Average number of concepts per document.
193    pub avg_concepts_per_doc: f64,
194}
195
196// ─────────────────────────────────────────────────────────────────────────────
197// FNV-1a helper
198// ─────────────────────────────────────────────────────────────────────────────
199
200/// Compute a 16-hex-char FNV-1a hash of the given string.
201fn fnv1a_str(s: &str) -> String {
202    let mut h: u64 = 14_695_981_039_346_656_037;
203    for b in s.bytes() {
204        h ^= b as u64;
205        h = h.wrapping_mul(1_099_511_628_211);
206    }
207    format!("{:016x}", h)
208}
209
210// ─────────────────────────────────────────────────────────────────────────────
211// KnowledgeBaseBuilder
212// ─────────────────────────────────────────────────────────────────────────────
213
214/// Incrementally builds and maintains a semantic knowledge base.
215///
216/// Supports entities, directed semantic relations, concept co-occurrence
217/// graphs, and document ingestion.
218#[derive(Debug, Default)]
219pub struct KnowledgeBaseBuilder {
220    /// All entities, keyed by entity ID.
221    entities: HashMap<String, KbBuilderEntity>,
222    /// All relations in insertion order.
223    relations: Vec<KbRelation>,
224    /// Concept co-occurrence graph, keyed by concept string.
225    concept_graph: HashMap<String, KbConceptNode>,
226    /// All ingested documents, keyed by doc_id.
227    documents: HashMap<String, KbDocument>,
228    /// Alias → canonical entity ID index (lower-cased alias).
229    entity_alias_index: HashMap<String, String>,
230}
231
232impl KnowledgeBaseBuilder {
233    // ── Construction ──────────────────────────────────────────────────────────
234
235    /// Create an empty `KnowledgeBaseBuilder`.
236    pub fn new() -> Self {
237        Self::default()
238    }
239
240    // ── Entity management ─────────────────────────────────────────────────────
241
242    /// Add a new entity to the knowledge base.
243    ///
244    /// Returns [`KbError::EntityAlreadyExists`] if an entity with the same ID
245    /// was previously registered.  Also registers all aliases in the alias index.
246    pub fn add_entity(&mut self, entity: KbBuilderEntity) -> Result<(), KbError> {
247        if self.entities.contains_key(&entity.id) {
248            return Err(KbError::EntityAlreadyExists(entity.id.clone()));
249        }
250
251        // Index the primary name as an alias.
252        let name_key = entity.name.to_lowercase();
253        self.entity_alias_index
254            .entry(name_key)
255            .or_insert_with(|| entity.id.clone());
256
257        // Index all explicit aliases.
258        for alias in &entity.aliases {
259            let alias_key = alias.to_lowercase();
260            self.entity_alias_index
261                .entry(alias_key)
262                .or_insert_with(|| entity.id.clone());
263        }
264
265        self.entities.insert(entity.id.clone(), entity);
266        Ok(())
267    }
268
269    /// Apply an in-place mutation to an existing entity.
270    ///
271    /// Returns [`KbError::EntityNotFound`] if the ID is not registered.
272    /// After mutation, re-indexes all aliases (including the primary name).
273    pub fn update_entity(
274        &mut self,
275        entity_id: &str,
276        update_fn: impl FnOnce(&mut KbBuilderEntity),
277    ) -> Result<(), KbError> {
278        let entity = self
279            .entities
280            .get_mut(entity_id)
281            .ok_or_else(|| KbError::EntityNotFound(entity_id.to_owned()))?;
282
283        // Remove stale alias entries for this entity before mutation.
284        let old_name = entity.name.to_lowercase();
285        let old_aliases: Vec<String> = entity.aliases.iter().map(|a| a.to_lowercase()).collect();
286
287        update_fn(entity);
288
289        // Update updated_at is caller responsibility via the closure; mark here too.
290        // Re-index aliases post mutation.
291        let new_name = entity.name.to_lowercase();
292        let new_aliases: Vec<String> = entity.aliases.iter().map(|a| a.to_lowercase()).collect();
293        let entity_id_owned = entity.id.clone();
294
295        // Remove old name/aliases.
296        if self
297            .entity_alias_index
298            .get(&old_name)
299            .map(|v| v == &entity_id_owned)
300            .unwrap_or(false)
301        {
302            self.entity_alias_index.remove(&old_name);
303        }
304        for old_alias in &old_aliases {
305            if self
306                .entity_alias_index
307                .get(old_alias)
308                .map(|v| v == &entity_id_owned)
309                .unwrap_or(false)
310            {
311                self.entity_alias_index.remove(old_alias);
312            }
313        }
314
315        // Insert new name/aliases.
316        self.entity_alias_index
317            .entry(new_name)
318            .or_insert(entity_id_owned.clone());
319        for new_alias in new_aliases {
320            self.entity_alias_index
321                .entry(new_alias)
322                .or_insert(entity_id_owned.clone());
323        }
324
325        Ok(())
326    }
327
328    /// Remove an entity and all relations that involve it.
329    ///
330    /// Returns `true` if the entity existed.
331    pub fn remove_entity(&mut self, entity_id: &str) -> bool {
332        let Some(entity) = self.entities.remove(entity_id) else {
333            return false;
334        };
335
336        // Remove alias entries that point to this entity.
337        let name_key = entity.name.to_lowercase();
338        if self
339            .entity_alias_index
340            .get(&name_key)
341            .map(|v| v == entity_id)
342            .unwrap_or(false)
343        {
344            self.entity_alias_index.remove(&name_key);
345        }
346        for alias in &entity.aliases {
347            let alias_key = alias.to_lowercase();
348            if self
349                .entity_alias_index
350                .get(&alias_key)
351                .map(|v| v == entity_id)
352                .unwrap_or(false)
353            {
354                self.entity_alias_index.remove(&alias_key);
355            }
356        }
357
358        // Remove relations involving this entity.
359        self.relations
360            .retain(|r| r.subject_id != entity_id && r.object_id != entity_id);
361
362        true
363    }
364
365    // ── Relation management ───────────────────────────────────────────────────
366
367    /// Add a semantic relation from a triple.
368    ///
369    /// Validates that both subject and object entities exist.  The relation ID
370    /// is the FNV-1a hash of `subject_id + predicate + object_id`.
371    ///
372    /// Returns the relation ID on success.  Returns
373    /// [`KbError::RelationAlreadyExists`] if the same triple was already added.
374    pub fn add_relation(
375        &mut self,
376        triple: KbTriple,
377        confidence: f64,
378        source: String,
379        now: u64,
380    ) -> Result<String, KbError> {
381        if !self.entities.contains_key(&triple.subject) {
382            return Err(KbError::EntityNotFound(triple.subject.clone()));
383        }
384        if !self.entities.contains_key(&triple.object) {
385            return Err(KbError::EntityNotFound(triple.object.clone()));
386        }
387
388        let relation_key = format!("{}{}{}", triple.subject, triple.predicate, triple.object);
389        let relation_id = fnv1a_str(&relation_key);
390
391        if self.relations.iter().any(|r| r.id == relation_id) {
392            return Err(KbError::RelationAlreadyExists(relation_id));
393        }
394
395        self.relations.push(KbRelation {
396            id: relation_id.clone(),
397            subject_id: triple.subject,
398            predicate: triple.predicate,
399            object_id: triple.object,
400            confidence,
401            source,
402            created_at: now,
403        });
404
405        Ok(relation_id)
406    }
407
408    /// Remove the relation with the given ID.
409    ///
410    /// Returns `true` if it existed.
411    pub fn remove_relation(&mut self, relation_id: &str) -> bool {
412        let before = self.relations.len();
413        self.relations.retain(|r| r.id != relation_id);
414        self.relations.len() < before
415    }
416
417    // ── Document management ───────────────────────────────────────────────────
418
419    /// Ingest a document into the knowledge base.
420    ///
421    /// Updates the concept graph for each concept in the document:
422    /// increments frequency, records the document ID, and adds
423    /// co-occurrence edges to every other concept in the same document.
424    ///
425    /// Returns [`KbError::DocumentAlreadyExists`] if the doc_id is taken.
426    pub fn add_document(&mut self, doc: KbDocument) -> Result<(), KbError> {
427        if self.documents.contains_key(&doc.doc_id) {
428            return Err(KbError::DocumentAlreadyExists(doc.doc_id.clone()));
429        }
430
431        let doc_id = doc.doc_id.clone();
432        let concepts = doc.concepts.clone();
433
434        // Ensure all concept nodes exist.
435        for concept in &concepts {
436            self.concept_graph
437                .entry(concept.clone())
438                .or_insert_with(|| KbConceptNode::new(concept.clone()));
439        }
440
441        // Update each concept node.
442        for (i, concept) in concepts.iter().enumerate() {
443            let node = self
444                .concept_graph
445                .get_mut(concept)
446                .expect("node was just inserted");
447            node.frequency += 1;
448            node.documents.push(doc_id.clone());
449
450            // Co-occurrence: add all other concepts in this document.
451            for (j, other) in concepts.iter().enumerate() {
452                if i != j && !node.related_concepts.contains(other) {
453                    node.related_concepts.push(other.clone());
454                }
455            }
456        }
457
458        self.documents.insert(doc_id, doc);
459        Ok(())
460    }
461
462    // ── Lookup helpers ────────────────────────────────────────────────────────
463
464    /// Look up an entity by its primary name or any registered alias.
465    ///
466    /// The comparison is case-insensitive.
467    pub fn find_entity_by_name(&self, name: &str) -> Option<&KbBuilderEntity> {
468        let key = name.to_lowercase();
469        // Check alias index first (covers primary name too).
470        if let Some(entity_id) = self.entity_alias_index.get(&key) {
471            return self.entities.get(entity_id);
472        }
473        // Fall back to linear scan over entity names.
474        self.entities
475            .values()
476            .find(|e| e.name.to_lowercase() == key)
477    }
478
479    /// Look up an entity via the alias index (case-insensitive).
480    pub fn find_entity_by_alias(&self, alias: &str) -> Option<&KbBuilderEntity> {
481        let key = alias.to_lowercase();
482        let entity_id = self.entity_alias_index.get(&key)?;
483        self.entities.get(entity_id)
484    }
485
486    /// All relations where the entity is the subject **or** the object.
487    pub fn relations_for_entity(&self, entity_id: &str) -> Vec<&KbRelation> {
488        self.relations
489            .iter()
490            .filter(|r| r.subject_id == entity_id || r.object_id == entity_id)
491            .collect()
492    }
493
494    /// Relations where `entity_id` is the subject (outgoing).
495    pub fn outgoing_relations(&self, entity_id: &str) -> Vec<&KbRelation> {
496        self.relations
497            .iter()
498            .filter(|r| r.subject_id == entity_id)
499            .collect()
500    }
501
502    /// Relations where `entity_id` is the object (incoming).
503    pub fn incoming_relations(&self, entity_id: &str) -> Vec<&KbRelation> {
504        self.relations
505            .iter()
506            .filter(|r| r.object_id == entity_id)
507            .collect()
508    }
509
510    /// Entities directly connected to `entity_id` via any relation (neighbours).
511    pub fn entity_neighbors(&self, entity_id: &str) -> Vec<&KbBuilderEntity> {
512        let mut seen: HashSet<&str> = HashSet::new();
513        let mut result = Vec::new();
514
515        for r in &self.relations {
516            let neighbor_id: Option<&str> = if r.subject_id == entity_id {
517                Some(&r.object_id)
518            } else if r.object_id == entity_id {
519                Some(&r.subject_id)
520            } else {
521                None
522            };
523
524            if let Some(nid) = neighbor_id {
525                if seen.insert(nid) {
526                    if let Some(e) = self.entities.get(nid) {
527                        result.push(e);
528                    }
529                }
530            }
531        }
532
533        result
534    }
535
536    /// Find the shortest path between two entities via BFS over relations.
537    ///
538    /// Returns a sequence of entity IDs (inclusive of `from_id` and `to_id`),
539    /// or `None` if no path exists within `max_hops`.
540    pub fn path_between(&self, from_id: &str, to_id: &str, max_hops: usize) -> Option<Vec<String>> {
541        if from_id == to_id {
542            return Some(vec![from_id.to_owned()]);
543        }
544        if max_hops == 0 {
545            return None;
546        }
547
548        // BFS: (current_node, path_so_far)
549        let mut visited: HashSet<String> = HashSet::new();
550        let mut queue: VecDeque<Vec<String>> = VecDeque::new();
551
552        visited.insert(from_id.to_owned());
553        queue.push_back(vec![from_id.to_owned()]);
554
555        while let Some(path) = queue.pop_front() {
556            let current = path.last().expect("path is non-empty");
557            let hops = path.len() - 1;
558
559            if hops >= max_hops {
560                continue;
561            }
562
563            for r in &self.relations {
564                let neighbor: Option<&str> = if r.subject_id == *current {
565                    Some(&r.object_id)
566                } else if r.object_id == *current {
567                    Some(&r.subject_id)
568                } else {
569                    None
570                };
571
572                if let Some(nid) = neighbor {
573                    if nid == to_id {
574                        let mut found = path.clone();
575                        found.push(nid.to_owned());
576                        return Some(found);
577                    }
578                    if !visited.contains(nid) {
579                        visited.insert(nid.to_owned());
580                        let mut next_path = path.clone();
581                        next_path.push(nid.to_owned());
582                        queue.push_back(next_path);
583                    }
584                }
585            }
586        }
587
588        None
589    }
590
591    // ── Concept graph ─────────────────────────────────────────────────────────
592
593    /// Return the top `n` concept nodes by frequency, descending.
594    pub fn top_concepts(&self, n: usize) -> Vec<&KbConceptNode> {
595        let mut nodes: Vec<&KbConceptNode> = self.concept_graph.values().collect();
596        nodes.sort_by(|a, b| {
597            b.frequency
598                .cmp(&a.frequency)
599                .then(a.concept.cmp(&b.concept))
600        });
601        nodes.truncate(n);
602        nodes
603    }
604
605    /// Count how many documents mention both `concept_a` and `concept_b`.
606    pub fn concept_cooccurrence(&self, concept_a: &str, concept_b: &str) -> usize {
607        let docs_a: HashSet<&str> = self
608            .concept_graph
609            .get(concept_a)
610            .map(|n| n.documents.iter().map(|d| d.as_str()).collect())
611            .unwrap_or_default();
612
613        let docs_b: HashSet<&str> = self
614            .concept_graph
615            .get(concept_b)
616            .map(|n| n.documents.iter().map(|d| d.as_str()).collect())
617            .unwrap_or_default();
618
619        docs_a.intersection(&docs_b).count()
620    }
621
622    // ── Counts ────────────────────────────────────────────────────────────────
623
624    /// Number of entities in the knowledge base.
625    pub fn entity_count(&self) -> usize {
626        self.entities.len()
627    }
628
629    /// Number of relations in the knowledge base.
630    pub fn relation_count(&self) -> usize {
631        self.relations.len()
632    }
633
634    /// Number of documents ingested.
635    pub fn document_count(&self) -> usize {
636        self.documents.len()
637    }
638
639    // ── Statistics ────────────────────────────────────────────────────────────
640
641    /// Compute aggregate statistics for the knowledge base.
642    pub fn stats(&self) -> KbStats {
643        let entity_count = self.entities.len();
644        let relation_count = self.relations.len();
645        let document_count = self.documents.len();
646        let concept_count = self.concept_graph.len();
647
648        let avg_relations_per_entity = if entity_count == 0 {
649            0.0
650        } else {
651            // Count total relation endpoints (each relation touches two entities).
652            (relation_count as f64 * 2.0) / entity_count as f64
653        };
654
655        let total_concepts: usize = self.documents.values().map(|d| d.concepts.len()).sum();
656        let avg_concepts_per_doc = if document_count == 0 {
657            0.0
658        } else {
659            total_concepts as f64 / document_count as f64
660        };
661
662        KbStats {
663            entity_count,
664            relation_count,
665            document_count,
666            concept_count,
667            avg_relations_per_entity,
668            avg_concepts_per_doc,
669        }
670    }
671
672    // ── Read-only access to internal collections ──────────────────────────────
673
674    /// Direct read access to the entity map.
675    pub fn entities(&self) -> &HashMap<String, KbBuilderEntity> {
676        &self.entities
677    }
678
679    /// Direct read access to the relation list.
680    pub fn relations(&self) -> &[KbRelation] {
681        &self.relations
682    }
683
684    /// Direct read access to the concept graph.
685    pub fn concept_graph(&self) -> &HashMap<String, KbConceptNode> {
686        &self.concept_graph
687    }
688
689    /// Direct read access to the document map.
690    pub fn documents(&self) -> &HashMap<String, KbDocument> {
691        &self.documents
692    }
693}
694
695// ─────────────────────────────────────────────────────────────────────────────
696// Tests
697// ─────────────────────────────────────────────────────────────────────────────
698
699#[cfg(test)]
700mod tests {
701    use crate::knowledge_base_builder::{
702        fnv1a_str, KbBuilderEntity, KbDocument, KbError, KbTriple, KnowledgeBaseBuilder,
703    };
704
705    // ── helpers ───────────────────────────────────────────────────────────────
706
707    fn make_entity(id: &str, name: &str, now: u64) -> KbBuilderEntity {
708        KbBuilderEntity::new(id, name, "concept", now)
709    }
710
711    fn make_entity_with_aliases(
712        id: &str,
713        name: &str,
714        aliases: Vec<&str>,
715        now: u64,
716    ) -> KbBuilderEntity {
717        let mut e = make_entity(id, name, now);
718        e.aliases = aliases.iter().map(|s| s.to_string()).collect();
719        e
720    }
721
722    fn make_doc(doc_id: &str, concepts: Vec<&str>, entities: Vec<&str>) -> KbDocument {
723        KbDocument::new(
724            doc_id,
725            "content",
726            entities.iter().map(|s| s.to_string()).collect(),
727            concepts.iter().map(|s| s.to_string()).collect(),
728            1_000_000,
729        )
730    }
731
732    // ── FNV-1a ────────────────────────────────────────────────────────────────
733
734    #[test]
735    fn test_fnv1a_empty_string() {
736        let h = fnv1a_str("");
737        assert_eq!(h.len(), 16);
738    }
739
740    #[test]
741    fn test_fnv1a_deterministic() {
742        let h1 = fnv1a_str("hello");
743        let h2 = fnv1a_str("hello");
744        assert_eq!(h1, h2);
745    }
746
747    #[test]
748    fn test_fnv1a_distinct() {
749        let h1 = fnv1a_str("abcpredxyz");
750        let h2 = fnv1a_str("xyzpredabc");
751        assert_ne!(h1, h2);
752    }
753
754    #[test]
755    fn test_fnv1a_hex_format() {
756        let h = fnv1a_str("test");
757        assert!(h.chars().all(|c| c.is_ascii_hexdigit()));
758        assert_eq!(h.len(), 16);
759    }
760
761    // ── KnowledgeBaseBuilder::new ─────────────────────────────────────────────
762
763    #[test]
764    fn test_new_is_empty() {
765        let kb = KnowledgeBaseBuilder::new();
766        assert_eq!(kb.entity_count(), 0);
767        assert_eq!(kb.relation_count(), 0);
768        assert_eq!(kb.document_count(), 0);
769    }
770
771    // ── add_entity ────────────────────────────────────────────────────────────
772
773    #[test]
774    fn test_add_entity_success() {
775        let mut kb = KnowledgeBaseBuilder::new();
776        let e = make_entity("e1", "Alpha", 0);
777        assert!(kb.add_entity(e).is_ok());
778        assert_eq!(kb.entity_count(), 1);
779    }
780
781    #[test]
782    fn test_add_entity_duplicate_returns_error() {
783        let mut kb = KnowledgeBaseBuilder::new();
784        kb.add_entity(make_entity("e1", "Alpha", 0))
785            .expect("test: add entity e1 for duplicate check");
786        let err = kb.add_entity(make_entity("e1", "Alpha2", 0)).unwrap_err();
787        assert_eq!(err, KbError::EntityAlreadyExists("e1".to_owned()));
788    }
789
790    #[test]
791    fn test_add_entity_alias_indexed() {
792        let mut kb = KnowledgeBaseBuilder::new();
793        let e = make_entity_with_aliases("e1", "Alpha", vec!["A", "al"], 0);
794        kb.add_entity(e).expect("test: add entity with aliases");
795        assert!(kb.find_entity_by_alias("a").is_some());
796        assert!(kb.find_entity_by_alias("al").is_some());
797        assert!(kb.find_entity_by_alias("ALPHA").is_some());
798    }
799
800    // ── update_entity ─────────────────────────────────────────────────────────
801
802    #[test]
803    fn test_update_entity_success() {
804        let mut kb = KnowledgeBaseBuilder::new();
805        kb.add_entity(make_entity("e1", "Alpha", 0))
806            .expect("test: add entity e1 before update");
807        kb.update_entity("e1", |e| {
808            e.name = "Beta".to_owned();
809        })
810        .expect("test: update entity e1 name to Beta");
811        let e = kb
812            .find_entity_by_name("Beta")
813            .expect("test: find entity by name Beta after update");
814        assert_eq!(e.id, "e1");
815    }
816
817    #[test]
818    fn test_update_entity_not_found() {
819        let mut kb = KnowledgeBaseBuilder::new();
820        let err = kb
821            .update_entity("nonexistent", |e| e.name = "x".to_owned())
822            .unwrap_err();
823        assert_eq!(err, KbError::EntityNotFound("nonexistent".to_owned()));
824    }
825
826    #[test]
827    fn test_update_entity_alias_reindexed() {
828        let mut kb = KnowledgeBaseBuilder::new();
829        let mut e = make_entity("e1", "Alpha", 0);
830        e.aliases = vec!["OldAlias".to_owned()];
831        kb.add_entity(e).expect("test: add entity with OldAlias");
832
833        kb.update_entity("e1", |ent| {
834            ent.aliases = vec!["NewAlias".to_owned()];
835        })
836        .expect("test: update entity e1 aliases to NewAlias");
837
838        assert!(kb.find_entity_by_alias("newalias").is_some());
839    }
840
841    // ── remove_entity ─────────────────────────────────────────────────────────
842
843    #[test]
844    fn test_remove_entity_existing() {
845        let mut kb = KnowledgeBaseBuilder::new();
846        kb.add_entity(make_entity("e1", "Alpha", 0))
847            .expect("test: add entity e1 before remove");
848        assert!(kb.remove_entity("e1"));
849        assert_eq!(kb.entity_count(), 0);
850    }
851
852    #[test]
853    fn test_remove_entity_nonexistent() {
854        let mut kb = KnowledgeBaseBuilder::new();
855        assert!(!kb.remove_entity("ghost"));
856    }
857
858    #[test]
859    fn test_remove_entity_cascades_relations() {
860        let mut kb = KnowledgeBaseBuilder::new();
861        kb.add_entity(make_entity("e1", "A", 0))
862            .expect("test: add entity e1 for cascade test");
863        kb.add_entity(make_entity("e2", "B", 0))
864            .expect("test: add entity e2 for cascade test");
865        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "test".into(), 0)
866            .expect("test: add relation e1 knows e2 for cascade test");
867        assert_eq!(kb.relation_count(), 1);
868
869        kb.remove_entity("e1");
870        assert_eq!(kb.relation_count(), 0);
871    }
872
873    #[test]
874    fn test_remove_entity_cleans_alias_index() {
875        let mut kb = KnowledgeBaseBuilder::new();
876        let mut e = make_entity("e1", "Alpha", 0);
877        e.aliases = vec!["AL".to_owned()];
878        kb.add_entity(e).expect("test: add entity with AL alias");
879        kb.remove_entity("e1");
880        assert!(kb.find_entity_by_alias("al").is_none());
881        assert!(kb.find_entity_by_name("alpha").is_none());
882    }
883
884    // ── add_relation ──────────────────────────────────────────────────────────
885
886    #[test]
887    fn test_add_relation_success() {
888        let mut kb = KnowledgeBaseBuilder::new();
889        kb.add_entity(make_entity("e1", "A", 0))
890            .expect("test: add entity e1 for relation test");
891        kb.add_entity(make_entity("e2", "B", 0))
892            .expect("test: add entity e2 for relation test");
893        let id = kb
894            .add_relation(KbTriple::new("e1", "knows", "e2"), 0.9, "src".into(), 0)
895            .expect("test: add relation e1 knows e2");
896        assert_eq!(id.len(), 16);
897        assert_eq!(kb.relation_count(), 1);
898    }
899
900    #[test]
901    fn test_add_relation_subject_missing() {
902        let mut kb = KnowledgeBaseBuilder::new();
903        kb.add_entity(make_entity("e2", "B", 0))
904            .expect("test: add entity e2 for subject-missing test");
905        let err = kb
906            .add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "src".into(), 0)
907            .unwrap_err();
908        assert_eq!(err, KbError::EntityNotFound("e1".to_owned()));
909    }
910
911    #[test]
912    fn test_add_relation_object_missing() {
913        let mut kb = KnowledgeBaseBuilder::new();
914        kb.add_entity(make_entity("e1", "A", 0))
915            .expect("test: add entity e1 for object-missing test");
916        let err = kb
917            .add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "src".into(), 0)
918            .unwrap_err();
919        assert_eq!(err, KbError::EntityNotFound("e2".to_owned()));
920    }
921
922    #[test]
923    fn test_add_relation_duplicate_error() {
924        let mut kb = KnowledgeBaseBuilder::new();
925        kb.add_entity(make_entity("e1", "A", 0))
926            .expect("test: add entity e1 for duplicate relation test");
927        kb.add_entity(make_entity("e2", "B", 0))
928            .expect("test: add entity e2 for duplicate relation test");
929        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "src".into(), 0)
930            .expect("test: add first relation e1 knows e2");
931        let err = kb
932            .add_relation(KbTriple::new("e1", "knows", "e2"), 0.5, "src2".into(), 1)
933            .unwrap_err();
934        matches!(err, KbError::RelationAlreadyExists(_));
935    }
936
937    #[test]
938    fn test_add_relation_different_predicates_allowed() {
939        let mut kb = KnowledgeBaseBuilder::new();
940        kb.add_entity(make_entity("e1", "A", 0))
941            .expect("test: add entity e1 for multi-predicate test");
942        kb.add_entity(make_entity("e2", "B", 0))
943            .expect("test: add entity e2 for multi-predicate test");
944        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "src".into(), 0)
945            .expect("test: add relation e1 knows e2");
946        kb.add_relation(KbTriple::new("e1", "hates", "e2"), 0.3, "src".into(), 0)
947            .expect("test: add relation e1 hates e2");
948        assert_eq!(kb.relation_count(), 2);
949    }
950
951    // ── remove_relation ───────────────────────────────────────────────────────
952
953    #[test]
954    fn test_remove_relation_existing() {
955        let mut kb = KnowledgeBaseBuilder::new();
956        kb.add_entity(make_entity("e1", "A", 0))
957            .expect("test: add entity e1 for remove relation test");
958        kb.add_entity(make_entity("e2", "B", 0))
959            .expect("test: add entity e2 for remove relation test");
960        let rid = kb
961            .add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "src".into(), 0)
962            .expect("test: add relation for removal");
963        assert!(kb.remove_relation(&rid));
964        assert_eq!(kb.relation_count(), 0);
965    }
966
967    #[test]
968    fn test_remove_relation_nonexistent() {
969        let mut kb = KnowledgeBaseBuilder::new();
970        assert!(!kb.remove_relation("deadbeefdeadbeef"));
971    }
972
973    // ── add_document ──────────────────────────────────────────────────────────
974
975    #[test]
976    fn test_add_document_success() {
977        let mut kb = KnowledgeBaseBuilder::new();
978        kb.add_document(make_doc("d1", vec!["rust", "memory"], vec![]))
979            .expect("test: add document d1 with rust and memory");
980        assert_eq!(kb.document_count(), 1);
981    }
982
983    #[test]
984    fn test_add_document_duplicate_error() {
985        let mut kb = KnowledgeBaseBuilder::new();
986        kb.add_document(make_doc("d1", vec!["rust"], vec![]))
987            .expect("test: add document d1 first time");
988        let err = kb
989            .add_document(make_doc("d1", vec!["rust"], vec![]))
990            .unwrap_err();
991        assert_eq!(err, KbError::DocumentAlreadyExists("d1".to_owned()));
992    }
993
994    #[test]
995    fn test_add_document_updates_concept_graph() {
996        let mut kb = KnowledgeBaseBuilder::new();
997        kb.add_document(make_doc("d1", vec!["rust", "memory"], vec![]))
998            .expect("test: add document d1 with rust and memory concepts");
999        let node = kb
1000            .concept_graph()
1001            .get("rust")
1002            .expect("test: concept rust should exist in graph");
1003        assert_eq!(node.frequency, 1);
1004        assert!(node.related_concepts.contains(&"memory".to_owned()));
1005    }
1006
1007    #[test]
1008    fn test_add_document_accumulates_frequency() {
1009        let mut kb = KnowledgeBaseBuilder::new();
1010        kb.add_document(make_doc("d1", vec!["rust"], vec![]))
1011            .expect("test: add document d1 with rust");
1012        kb.add_document(make_doc("d2", vec!["rust"], vec![]))
1013            .expect("test: add document d2 with rust");
1014        let node = kb
1015            .concept_graph()
1016            .get("rust")
1017            .expect("test: concept rust should exist in graph after two docs");
1018        assert_eq!(node.frequency, 2);
1019    }
1020
1021    // ── find_entity_by_name / find_entity_by_alias ───────────────────────────
1022
1023    #[test]
1024    fn test_find_entity_by_name_exact() {
1025        let mut kb = KnowledgeBaseBuilder::new();
1026        kb.add_entity(make_entity("e1", "RustLang", 0))
1027            .expect("test: add entity RustLang for name lookup");
1028        let e = kb
1029            .find_entity_by_name("RustLang")
1030            .expect("test: find entity by exact name RustLang");
1031        assert_eq!(e.id, "e1");
1032    }
1033
1034    #[test]
1035    fn test_find_entity_by_name_case_insensitive() {
1036        let mut kb = KnowledgeBaseBuilder::new();
1037        kb.add_entity(make_entity("e1", "RustLang", 0))
1038            .expect("test: add entity RustLang for case-insensitive lookup");
1039        assert!(kb.find_entity_by_name("rustlang").is_some());
1040        assert!(kb.find_entity_by_name("RUSTLANG").is_some());
1041    }
1042
1043    #[test]
1044    fn test_find_entity_by_name_not_found() {
1045        let kb = KnowledgeBaseBuilder::new();
1046        assert!(kb.find_entity_by_name("nobody").is_none());
1047    }
1048
1049    #[test]
1050    fn test_find_entity_by_alias() {
1051        let mut kb = KnowledgeBaseBuilder::new();
1052        let e = make_entity_with_aliases("e1", "RustLang", vec!["rs", "rust"], 0);
1053        kb.add_entity(e)
1054            .expect("test: add entity with rs and rust aliases");
1055        assert!(kb.find_entity_by_alias("RS").is_some());
1056        assert!(kb.find_entity_by_alias("rust").is_some());
1057    }
1058
1059    // ── relation queries ──────────────────────────────────────────────────────
1060
1061    #[test]
1062    fn test_relations_for_entity() {
1063        let mut kb = KnowledgeBaseBuilder::new();
1064        kb.add_entity(make_entity("e1", "A", 0))
1065            .expect("test: add entity e1 for relations query");
1066        kb.add_entity(make_entity("e2", "B", 0))
1067            .expect("test: add entity e2 for relations query");
1068        kb.add_entity(make_entity("e3", "C", 0))
1069            .expect("test: add entity e3 for relations query");
1070        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1071            .expect("test: add relation e1 knows e2");
1072        kb.add_relation(KbTriple::new("e3", "mentions", "e1"), 0.8, "s".into(), 0)
1073            .expect("test: add relation e3 mentions e1");
1074        kb.add_relation(KbTriple::new("e2", "linked", "e3"), 0.5, "s".into(), 0)
1075            .expect("test: add relation e2 linked e3");
1076
1077        let rels = kb.relations_for_entity("e1");
1078        assert_eq!(rels.len(), 2);
1079    }
1080
1081    #[test]
1082    fn test_outgoing_relations() {
1083        let mut kb = KnowledgeBaseBuilder::new();
1084        kb.add_entity(make_entity("e1", "A", 0))
1085            .expect("test: add entity e1 for outgoing relations");
1086        kb.add_entity(make_entity("e2", "B", 0))
1087            .expect("test: add entity e2 for outgoing relations");
1088        kb.add_entity(make_entity("e3", "C", 0))
1089            .expect("test: add entity e3 for outgoing relations");
1090        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1091            .expect("test: add relation e1 knows e2 for outgoing");
1092        kb.add_relation(KbTriple::new("e3", "knows", "e1"), 0.5, "s".into(), 0)
1093            .expect("test: add relation e3 knows e1 for outgoing");
1094
1095        let out = kb.outgoing_relations("e1");
1096        assert_eq!(out.len(), 1);
1097        assert_eq!(out[0].object_id, "e2");
1098    }
1099
1100    #[test]
1101    fn test_incoming_relations() {
1102        let mut kb = KnowledgeBaseBuilder::new();
1103        kb.add_entity(make_entity("e1", "A", 0))
1104            .expect("test: add entity e1 for incoming relations");
1105        kb.add_entity(make_entity("e2", "B", 0))
1106            .expect("test: add entity e2 for incoming relations");
1107        kb.add_entity(make_entity("e3", "C", 0))
1108            .expect("test: add entity e3 for incoming relations");
1109        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1110            .expect("test: add relation e1 knows e2 for incoming");
1111        kb.add_relation(KbTriple::new("e3", "knows", "e2"), 0.5, "s".into(), 0)
1112            .expect("test: add relation e3 knows e2 for incoming");
1113
1114        let inc = kb.incoming_relations("e2");
1115        assert_eq!(inc.len(), 2);
1116    }
1117
1118    // ── entity_neighbors ─────────────────────────────────────────────────────
1119
1120    #[test]
1121    fn test_entity_neighbors() {
1122        let mut kb = KnowledgeBaseBuilder::new();
1123        kb.add_entity(make_entity("e1", "A", 0))
1124            .expect("test: add entity e1 for neighbors");
1125        kb.add_entity(make_entity("e2", "B", 0))
1126            .expect("test: add entity e2 for neighbors");
1127        kb.add_entity(make_entity("e3", "C", 0))
1128            .expect("test: add entity e3 for neighbors");
1129        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1130            .expect("test: add relation e1 knows e2 for neighbors");
1131        kb.add_relation(KbTriple::new("e3", "knows", "e1"), 0.8, "s".into(), 0)
1132            .expect("test: add relation e3 knows e1 for neighbors");
1133
1134        let neighbors = kb.entity_neighbors("e1");
1135        let ids: Vec<&str> = neighbors.iter().map(|e| e.id.as_str()).collect();
1136        assert!(ids.contains(&"e2"));
1137        assert!(ids.contains(&"e3"));
1138    }
1139
1140    #[test]
1141    fn test_entity_neighbors_no_duplicates() {
1142        let mut kb = KnowledgeBaseBuilder::new();
1143        kb.add_entity(make_entity("e1", "A", 0))
1144            .expect("test: add entity e1 for no-duplicate neighbors");
1145        kb.add_entity(make_entity("e2", "B", 0))
1146            .expect("test: add entity e2 for no-duplicate neighbors");
1147        // Two different predicates between the same pair.
1148        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1149            .expect("test: add relation e1 knows e2 for no-duplicate");
1150        kb.add_relation(KbTriple::new("e1", "likes", "e2"), 0.5, "s".into(), 0)
1151            .expect("test: add relation e1 likes e2 for no-duplicate");
1152
1153        let neighbors = kb.entity_neighbors("e1");
1154        assert_eq!(neighbors.len(), 1, "neighbor e2 should appear once");
1155    }
1156
1157    // ── path_between ──────────────────────────────────────────────────────────
1158
1159    #[test]
1160    fn test_path_between_same_node() {
1161        let mut kb = KnowledgeBaseBuilder::new();
1162        kb.add_entity(make_entity("e1", "A", 0))
1163            .expect("test: add entity e1 for same-node path");
1164        let path = kb
1165            .path_between("e1", "e1", 5)
1166            .expect("test: path from e1 to e1 should succeed");
1167        assert_eq!(path, vec!["e1"]);
1168    }
1169
1170    #[test]
1171    fn test_path_between_direct() {
1172        let mut kb = KnowledgeBaseBuilder::new();
1173        kb.add_entity(make_entity("e1", "A", 0))
1174            .expect("test: add entity e1 for direct path");
1175        kb.add_entity(make_entity("e2", "B", 0))
1176            .expect("test: add entity e2 for direct path");
1177        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1178            .expect("test: add relation e1 knows e2 for direct path");
1179
1180        let path = kb
1181            .path_between("e1", "e2", 5)
1182            .expect("test: path between e1 and e2 should succeed");
1183        assert_eq!(path, vec!["e1", "e2"]);
1184    }
1185
1186    #[test]
1187    fn test_path_between_two_hops() {
1188        let mut kb = KnowledgeBaseBuilder::new();
1189        kb.add_entity(make_entity("e1", "A", 0))
1190            .expect("test: add entity e1 for two-hop path");
1191        kb.add_entity(make_entity("e2", "B", 0))
1192            .expect("test: add entity e2 for two-hop path");
1193        kb.add_entity(make_entity("e3", "C", 0))
1194            .expect("test: add entity e3 for two-hop path");
1195        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1196            .expect("test: add relation e1 knows e2 for two-hop");
1197        kb.add_relation(KbTriple::new("e2", "knows", "e3"), 1.0, "s".into(), 0)
1198            .expect("test: add relation e2 knows e3 for two-hop");
1199
1200        let path = kb
1201            .path_between("e1", "e3", 5)
1202            .expect("test: two-hop path from e1 to e3 should succeed");
1203        assert_eq!(path.first().map(|s| s.as_str()), Some("e1"));
1204        assert_eq!(path.last().map(|s| s.as_str()), Some("e3"));
1205        assert_eq!(path.len(), 3);
1206    }
1207
1208    #[test]
1209    fn test_path_between_no_path() {
1210        let mut kb = KnowledgeBaseBuilder::new();
1211        kb.add_entity(make_entity("e1", "A", 0))
1212            .expect("test: add entity e1 for no-path test");
1213        kb.add_entity(make_entity("e2", "B", 0))
1214            .expect("test: add entity e2 for no-path test");
1215        // No relation between them.
1216        assert!(kb.path_between("e1", "e2", 5).is_none());
1217    }
1218
1219    #[test]
1220    fn test_path_between_exceeds_max_hops() {
1221        let mut kb = KnowledgeBaseBuilder::new();
1222        for i in 0..5u8 {
1223            kb.add_entity(make_entity(&format!("e{}", i), &format!("E{}", i), 0))
1224                .expect("test: add entity in path chain");
1225        }
1226        for i in 0..4u8 {
1227            let s = format!("e{}", i);
1228            let o = format!("e{}", i + 1);
1229            kb.add_relation(KbTriple::new(s, "linked", o), 1.0, "s".into(), 0)
1230                .expect("test: add relation in path chain");
1231        }
1232        // Path e0 → e4 requires 4 hops; limit to 2.
1233        assert!(kb.path_between("e0", "e4", 2).is_none());
1234        // Limit to 5 should work.
1235        assert!(kb.path_between("e0", "e4", 5).is_some());
1236    }
1237
1238    // ── top_concepts ──────────────────────────────────────────────────────────
1239
1240    #[test]
1241    fn test_top_concepts_ordering() {
1242        let mut kb = KnowledgeBaseBuilder::new();
1243        kb.add_document(make_doc("d1", vec!["rust"], vec![]))
1244            .expect("test: add document d1 with rust");
1245        kb.add_document(make_doc("d2", vec!["rust", "memory"], vec![]))
1246            .expect("test: add document d2 with rust and memory");
1247        kb.add_document(make_doc("d3", vec!["memory", "safety"], vec![]))
1248            .expect("test: add document d3 with memory and safety");
1249
1250        let top = kb.top_concepts(2);
1251        assert_eq!(top.len(), 2);
1252        // "rust" and "memory" both have frequency 2; "safety" has 1.
1253        let freqs: Vec<u32> = top.iter().map(|n| n.frequency).collect();
1254        assert!(freqs[0] >= freqs[1]);
1255    }
1256
1257    #[test]
1258    fn test_top_concepts_fewer_than_n() {
1259        let mut kb = KnowledgeBaseBuilder::new();
1260        kb.add_document(make_doc("d1", vec!["rust"], vec![]))
1261            .expect("test: add document d1 for concept count check");
1262        let top = kb.top_concepts(10);
1263        assert_eq!(top.len(), 1);
1264    }
1265
1266    // ── concept_cooccurrence ─────────────────────────────────────────────────
1267
1268    #[test]
1269    fn test_concept_cooccurrence_positive() {
1270        let mut kb = KnowledgeBaseBuilder::new();
1271        kb.add_document(make_doc("d1", vec!["rust", "memory"], vec![]))
1272            .expect("test: add document d1 with rust and memory");
1273        kb.add_document(make_doc("d2", vec!["rust", "safety"], vec![]))
1274            .expect("test: add document d2 with rust and safety");
1275        kb.add_document(make_doc("d3", vec!["memory", "safety"], vec![]))
1276            .expect("test: add document d3 with memory and safety");
1277
1278        assert_eq!(kb.concept_cooccurrence("rust", "memory"), 1);
1279        assert_eq!(kb.concept_cooccurrence("rust", "safety"), 1);
1280        assert_eq!(kb.concept_cooccurrence("memory", "safety"), 1);
1281    }
1282
1283    #[test]
1284    fn test_concept_cooccurrence_zero() {
1285        let mut kb = KnowledgeBaseBuilder::new();
1286        kb.add_document(make_doc("d1", vec!["rust"], vec![]))
1287            .expect("test: add document d1 with rust");
1288        kb.add_document(make_doc("d2", vec!["python"], vec![]))
1289            .expect("test: add document d2 with python");
1290        assert_eq!(kb.concept_cooccurrence("rust", "python"), 0);
1291    }
1292
1293    #[test]
1294    fn test_concept_cooccurrence_unknown_concept() {
1295        let kb = KnowledgeBaseBuilder::new();
1296        assert_eq!(kb.concept_cooccurrence("unknown_a", "unknown_b"), 0);
1297    }
1298
1299    // ── stats ─────────────────────────────────────────────────────────────────
1300
1301    #[test]
1302    fn test_stats_empty() {
1303        let kb = KnowledgeBaseBuilder::new();
1304        let s = kb.stats();
1305        assert_eq!(s.entity_count, 0);
1306        assert_eq!(s.relation_count, 0);
1307        assert_eq!(s.document_count, 0);
1308        assert_eq!(s.concept_count, 0);
1309        assert!((s.avg_relations_per_entity - 0.0).abs() < 1e-9);
1310        assert!((s.avg_concepts_per_doc - 0.0).abs() < 1e-9);
1311    }
1312
1313    #[test]
1314    fn test_stats_with_data() {
1315        let mut kb = KnowledgeBaseBuilder::new();
1316        kb.add_entity(make_entity("e1", "A", 0))
1317            .expect("test: add entity e1 for stats");
1318        kb.add_entity(make_entity("e2", "B", 0))
1319            .expect("test: add entity e2 for stats");
1320        kb.add_relation(KbTriple::new("e1", "knows", "e2"), 1.0, "s".into(), 0)
1321            .expect("test: add relation e1 knows e2 for stats");
1322        kb.add_document(make_doc("d1", vec!["c1", "c2", "c3"], vec![]))
1323            .expect("test: add document d1 for stats");
1324
1325        let s = kb.stats();
1326        assert_eq!(s.entity_count, 2);
1327        assert_eq!(s.relation_count, 1);
1328        assert_eq!(s.document_count, 1);
1329        assert_eq!(s.concept_count, 3);
1330        // 1 relation * 2 endpoints / 2 entities = 1.0
1331        assert!((s.avg_relations_per_entity - 1.0).abs() < 1e-9);
1332        // 3 concepts / 1 document = 3.0
1333        assert!((s.avg_concepts_per_doc - 3.0).abs() < 1e-9);
1334    }
1335
1336    // ── KbBuilderEntity::new ──────────────────────────────────────────────────
1337
1338    #[test]
1339    fn test_kb_builder_entity_new() {
1340        let e = KbBuilderEntity::new("id1", "Name", "person", 42);
1341        assert_eq!(e.id, "id1");
1342        assert_eq!(e.name, "Name");
1343        assert_eq!(e.entity_type, "person");
1344        assert_eq!(e.created_at, 42);
1345        assert_eq!(e.updated_at, 42);
1346        assert!(e.aliases.is_empty());
1347        assert!(e.embedding.is_none());
1348    }
1349
1350    // ── KbTriple::new ─────────────────────────────────────────────────────────
1351
1352    #[test]
1353    fn test_kb_triple_new() {
1354        let t = KbTriple::new("s", "p", "o");
1355        assert_eq!(t.subject, "s");
1356        assert_eq!(t.predicate, "p");
1357        assert_eq!(t.object, "o");
1358    }
1359
1360    // ── KbDocument::new ───────────────────────────────────────────────────────
1361
1362    #[test]
1363    fn test_kb_document_new() {
1364        let d = KbDocument::new("d1", "hello", vec!["e1".into()], vec!["c1".into()], 100);
1365        assert_eq!(d.doc_id, "d1");
1366        assert_eq!(d.content, "hello");
1367        assert_eq!(d.entities_mentioned, vec!["e1"]);
1368        assert_eq!(d.concepts, vec!["c1"]);
1369        assert_eq!(d.added_at, 100);
1370    }
1371
1372    // ── integration ───────────────────────────────────────────────────────────
1373
1374    #[test]
1375    fn test_integration_full_graph() {
1376        let mut kb = KnowledgeBaseBuilder::new();
1377
1378        // Build a small knowledge base.
1379        for (id, name) in [("alice", "Alice"), ("bob", "Bob"), ("carol", "Carol")] {
1380            kb.add_entity(KbBuilderEntity::new(id, name, "person", 0))
1381                .expect("test: add person entity for integration test");
1382        }
1383        kb.add_relation(KbTriple::new("alice", "knows", "bob"), 0.9, "src".into(), 0)
1384            .expect("test: add relation alice knows bob");
1385        kb.add_relation(KbTriple::new("bob", "knows", "carol"), 0.8, "src".into(), 0)
1386            .expect("test: add relation bob knows carol");
1387
1388        let path = kb
1389            .path_between("alice", "carol", 3)
1390            .expect("test: path from alice to carol through bob");
1391        assert_eq!(path, vec!["alice", "bob", "carol"]);
1392
1393        let s = kb.stats();
1394        assert_eq!(s.entity_count, 3);
1395        assert_eq!(s.relation_count, 2);
1396    }
1397
1398    #[test]
1399    fn test_integration_document_entity_linking() {
1400        let mut kb = KnowledgeBaseBuilder::new();
1401        kb.add_entity(make_entity("e1", "Rust", 0))
1402            .expect("test: add entity Rust for document linking");
1403        let doc = KbDocument::new(
1404            "d1",
1405            "Rust is a systems language",
1406            vec!["e1".to_owned()],
1407            vec!["systems".to_owned(), "language".to_owned()],
1408            0,
1409        );
1410        kb.add_document(doc)
1411            .expect("test: add document d1 linking to Rust entity");
1412
1413        let node = kb
1414            .concept_graph()
1415            .get("systems")
1416            .expect("test: concept systems should exist in graph");
1417        assert!(node.documents.contains(&"d1".to_owned()));
1418        assert!(node.related_concepts.contains(&"language".to_owned()));
1419    }
1420}