omnivore_core/graph/
builder.rs1use 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}