omnivore_core/graph/
builder.rs

1use crate::graph::{Edge, KnowledgeGraph, Node};
2use crate::Result;
3use std::collections::HashMap;
4
5pub struct GraphBuilder {
6    graph: KnowledgeGraph,
7}
8
9impl Default for GraphBuilder {
10    fn default() -> Self {
11        Self::new()
12    }
13}
14
15impl GraphBuilder {
16    pub fn new() -> Self {
17        Self {
18            graph: KnowledgeGraph::new(),
19        }
20    }
21
22    pub fn add_entity(
23        &mut self,
24        id: String,
25        entity_type: String,
26        properties: HashMap<String, serde_json::Value>,
27    ) -> Result<()> {
28        let node = Node {
29            id,
30            node_type: entity_type,
31            properties,
32        };
33        self.graph.add_node(node)
34    }
35
36    pub fn add_relationship(
37        &mut self,
38        from: String,
39        to: String,
40        rel_type: String,
41        properties: HashMap<String, serde_json::Value>,
42    ) -> Result<()> {
43        let edge = Edge {
44            from,
45            to,
46            edge_type: rel_type,
47            properties,
48        };
49        self.graph.add_edge(edge)
50    }
51
52    pub fn build(self) -> KnowledgeGraph {
53        self.graph
54    }
55}