Skip to main content

entrenar/monitor/inference/provenance/
graph.rs

1//! Provenance graph for causal analysis.
2
3use super::edge::ProvenanceEdge;
4use super::node::{NodeId, ProvenanceNode};
5use crate::monitor::inference::path::DecisionPath;
6use crate::monitor::inference::trace::DecisionTrace;
7use std::collections::HashMap;
8
9/// Provenance graph for causal analysis
10///
11/// # Features
12/// - DAG structure for causal relationships
13/// - Backward traversal for root cause analysis
14/// - Anomaly detection in decision chains
15pub struct ProvenanceGraph {
16    /// Nodes indexed by ID
17    nodes: HashMap<NodeId, ProvenanceNode>,
18    /// Edges
19    edges: Vec<ProvenanceEdge>,
20    /// Adjacency list (forward): node -> outgoing edges
21    forward: HashMap<NodeId, Vec<usize>>,
22    /// Adjacency list (backward): node -> incoming edges
23    backward: HashMap<NodeId, Vec<usize>>,
24    /// Next node ID
25    next_id: NodeId,
26}
27
28impl ProvenanceGraph {
29    /// Create a new empty graph
30    pub fn new() -> Self {
31        Self {
32            nodes: HashMap::new(),
33            edges: Vec::new(),
34            forward: HashMap::new(),
35            backward: HashMap::new(),
36            next_id: 0,
37        }
38    }
39
40    /// Add a node to the graph
41    pub fn add_node(&mut self, node: ProvenanceNode) -> NodeId {
42        let id = self.next_id;
43        self.next_id += 1;
44        self.nodes.insert(id, node);
45        id
46    }
47
48    /// Add an edge to the graph
49    pub fn add_edge(&mut self, edge: ProvenanceEdge) {
50        let edge_idx = self.edges.len();
51        self.forward.entry(edge.from).or_default().push(edge_idx);
52        self.backward.entry(edge.to).or_default().push(edge_idx);
53        self.edges.push(edge);
54    }
55
56    /// Get a node by ID
57    pub fn get_node(&self, id: NodeId) -> Option<&ProvenanceNode> {
58        self.nodes.get(&id)
59    }
60
61    /// Get all nodes
62    pub fn nodes(&self) -> &HashMap<NodeId, ProvenanceNode> {
63        &self.nodes
64    }
65
66    /// Get all edges
67    pub fn edges(&self) -> &[ProvenanceEdge] {
68        &self.edges
69    }
70
71    /// Get incoming edges for a node
72    pub fn incoming_edges(&self, id: NodeId) -> Vec<&ProvenanceEdge> {
73        self.backward
74            .get(&id)
75            .map(|indices| indices.iter().map(|&i| &self.edges[i]).collect())
76            .unwrap_or_default()
77    }
78
79    /// Get outgoing edges for a node
80    pub fn outgoing_edges(&self, id: NodeId) -> Vec<&ProvenanceEdge> {
81        self.forward
82            .get(&id)
83            .map(|indices| indices.iter().map(|&i| &self.edges[i]).collect())
84            .unwrap_or_default()
85    }
86
87    /// Get predecessor nodes
88    pub fn predecessors(&self, id: NodeId) -> Vec<NodeId> {
89        self.incoming_edges(id).into_iter().map(|e| e.from).collect()
90    }
91
92    /// Get successor nodes
93    pub fn successors(&self, id: NodeId) -> Vec<NodeId> {
94        self.outgoing_edges(id).into_iter().map(|e| e.to).collect()
95    }
96
97    /// Number of nodes
98    pub fn node_count(&self) -> usize {
99        self.nodes.len()
100    }
101
102    /// Number of edges
103    pub fn edge_count(&self) -> usize {
104        self.edges.len()
105    }
106
107    /// Add an inference node from a decision trace
108    pub fn add_inference<P: DecisionPath>(
109        &mut self,
110        trace: &DecisionTrace<P>,
111        model_id: &str,
112        model_version: &str,
113    ) -> NodeId {
114        self.add_node(ProvenanceNode::Inference {
115            model_id: model_id.to_string(),
116            model_version: model_version.to_string(),
117            confidence: trace.confidence(),
118            output: trace.output,
119        })
120    }
121}
122
123impl Default for ProvenanceGraph {
124    fn default() -> Self {
125        Self::new()
126    }
127}