kopitiam_knowledge/
graph.rs

1use std::collections::HashMap;
2
3use kopitiam_ontology::{Entity, EntityId, EntityKind, Relationship, RelationshipId};
4use serde::{Deserialize, Serialize};
5
6/// The unified knowledge graph: every entity and relationship produced by
7/// any [`kopitiam_ontology`]-speaking provider, merged into one structure.
8///
9/// This type is storage-agnostic on purpose. It holds no file handles and
10/// makes no assumptions about where it is persisted — it derives
11/// `Serialize`/`Deserialize` so `kopitiam-index` (or anything else) can snapshot
12/// it without `kopitiam-knowledge` depending on a storage engine.
13#[derive(Debug, Default, Clone, Serialize, Deserialize)]
14pub struct SemanticGraph {
15    entities: HashMap<EntityId, Entity>,
16    relationships: HashMap<RelationshipId, Relationship>,
17    #[serde(default)]
18    outgoing: HashMap<EntityId, Vec<RelationshipId>>,
19    #[serde(default)]
20    incoming: HashMap<EntityId, Vec<RelationshipId>>,
21}
22
23impl SemanticGraph {
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    /// Inserts one entity, keyed by its own [`EntityId`]. Re-inserting an
29    /// id already present overwrites the previous entity.
30    pub fn insert_entity(&mut self, entity: Entity) -> EntityId {
31        let id = entity.id;
32        self.entities.insert(id, entity);
33        id
34    }
35
36    /// Inserts one relationship and indexes it for traversal in both
37    /// directions.
38    pub fn insert_relationship(&mut self, relationship: Relationship) -> RelationshipId {
39        let id = relationship.id;
40        self.outgoing.entry(relationship.from).or_default().push(id);
41        self.incoming.entry(relationship.to).or_default().push(id);
42        self.relationships.insert(id, relationship);
43        id
44    }
45
46    /// Merges a batch of facts from a provider run in one call — the usual
47    /// entry point after calling a `kopitiam_semantic::KnowledgeProvider`.
48    pub fn extend(
49        &mut self,
50        entities: impl IntoIterator<Item = Entity>,
51        relationships: impl IntoIterator<Item = Relationship>,
52    ) {
53        for entity in entities {
54            self.insert_entity(entity);
55        }
56        for relationship in relationships {
57            self.insert_relationship(relationship);
58        }
59    }
60
61    pub fn entity(&self, id: EntityId) -> Option<&Entity> {
62        self.entities.get(&id)
63    }
64
65    pub fn relationship(&self, id: RelationshipId) -> Option<&Relationship> {
66        self.relationships.get(&id)
67    }
68
69    pub fn entity_count(&self) -> usize {
70        self.entities.len()
71    }
72
73    pub fn relationship_count(&self) -> usize {
74        self.relationships.len()
75    }
76
77    pub fn entities(&self) -> impl Iterator<Item = &Entity> {
78        self.entities.values()
79    }
80
81    pub fn entities_of_kind(&self, kind: EntityKind) -> impl Iterator<Item = &Entity> {
82        self.entities.values().filter(move |e| e.kind == kind)
83    }
84
85    /// Relationships whose `from` is `id`.
86    pub fn relationships_from(&self, id: EntityId) -> impl Iterator<Item = &Relationship> {
87        self.outgoing
88            .get(&id)
89            .into_iter()
90            .flatten()
91            .filter_map(move |rel_id| self.relationships.get(rel_id))
92    }
93
94    /// Relationships whose `to` is `id`.
95    pub fn relationships_to(&self, id: EntityId) -> impl Iterator<Item = &Relationship> {
96        self.incoming
97            .get(&id)
98            .into_iter()
99            .flatten()
100            .filter_map(move |rel_id| self.relationships.get(rel_id))
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use kopitiam_ontology::RelationshipKind;
108
109    #[test]
110    fn inserts_and_queries_entities() {
111        let mut graph = SemanticGraph::new();
112        let a = Entity::new(EntityKind::Artifact, "kopitiam-knowledge", "test");
113        let b = Entity::new(EntityKind::Symbol, "SemanticGraph", "test");
114        let (a_id, b_id) = (a.id, b.id);
115        graph.extend([a, b], []);
116
117        assert_eq!(graph.entity_count(), 2);
118        assert_eq!(graph.entities_of_kind(EntityKind::Symbol).count(), 1);
119        assert!(graph.entity(a_id).is_some());
120        assert!(graph.entity(b_id).is_some());
121    }
122
123    #[test]
124    fn traverses_relationships_in_both_directions() {
125        let mut graph = SemanticGraph::new();
126        let a = Entity::new(EntityKind::Symbol, "SemanticGraph", "test");
127        let b = Entity::new(EntityKind::Artifact, "kopitiam-knowledge", "test");
128        let (a_id, b_id) = (a.id, b.id);
129        graph.insert_entity(a);
130        graph.insert_entity(b);
131        graph.insert_relationship(Relationship::new(a_id, b_id, RelationshipKind::LocatedIn));
132
133        let out: Vec<_> = graph.relationships_from(a_id).collect();
134        assert_eq!(out.len(), 1);
135        assert_eq!(out[0].kind, RelationshipKind::LocatedIn);
136
137        let inn: Vec<_> = graph.relationships_to(b_id).collect();
138        assert_eq!(inn.len(), 1);
139        assert_eq!(inn[0].from, a_id);
140    }
141
142    #[test]
143    fn round_trips_through_json() {
144        let mut graph = SemanticGraph::new();
145        let a = Entity::new(EntityKind::Artifact, "a", "test");
146        let b = Entity::new(EntityKind::Symbol, "b", "test");
147        let (a_id, b_id) = (a.id, b.id);
148        graph.extend([a, b], [Relationship::new(a_id, b_id, RelationshipKind::DependsOn)]);
149
150        let json = serde_json::to_string(&graph).unwrap();
151        let back: SemanticGraph = serde_json::from_str(&json).unwrap();
152        assert_eq!(back.entity_count(), 2);
153        assert_eq!(back.relationship_count(), 1);
154        assert_eq!(back.relationships_from(a_id).count(), 1);
155    }
156}
kopitiam_knowledge/graph.rs

kopitiam_knowledge/
graph.rs
