Skip to main content

brainwires_core/
graph.rs

1//! Graph types and traits for knowledge graph abstraction.
2//!
3//! Defines entity types, edge types, and trait interfaces for entity stores
4//! and relationship graphs. These abstractions allow consumers (like SEAL)
5//! to depend on traits rather than concrete storage implementations.
6
7use serde::{Deserialize, Serialize};
8
9/// Types of entities tracked in the knowledge graph.
10#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
11pub enum EntityType {
12    /// A source file entity.
13    File,
14    /// A function or method entity.
15    Function,
16    /// A type (struct, enum, class) entity.
17    Type,
18    /// A variable or constant entity.
19    Variable,
20    /// An abstract concept entity.
21    Concept,
22    /// An error or exception entity.
23    Error,
24    /// A CLI command entity.
25    Command,
26}
27
28impl EntityType {
29    /// Returns the string representation of this entity type.
30    pub fn as_str(&self) -> &'static str {
31        match self {
32            EntityType::File => "file",
33            EntityType::Function => "function",
34            EntityType::Type => "type",
35            EntityType::Variable => "variable",
36            EntityType::Concept => "concept",
37            EntityType::Error => "error",
38            EntityType::Command => "command",
39        }
40    }
41}
42
43/// Types of edges in the relationship graph.
44#[derive(Debug, Clone, PartialEq, Eq, Hash)]
45pub enum EdgeType {
46    /// Entities co-occur in the same context.
47    CoOccurs,
48    /// One entity contains another.
49    Contains,
50    /// One entity references another.
51    References,
52    /// One entity depends on another.
53    DependsOn,
54    /// One entity modifies another.
55    Modifies,
56    /// One entity defines another.
57    Defines,
58}
59
60impl EdgeType {
61    /// Get the default weight for this edge type.
62    pub fn weight(&self) -> f32 {
63        match self {
64            EdgeType::Defines => 1.0,
65            EdgeType::Contains => 0.9,
66            EdgeType::DependsOn => 0.8,
67            EdgeType::Modifies => 0.7,
68            EdgeType::References => 0.6,
69            EdgeType::CoOccurs => 0.3,
70        }
71    }
72}
73
74/// A node in the relationship graph.
75#[derive(Debug, Clone)]
76pub struct GraphNode {
77    /// Name of the entity.
78    pub entity_name: String,
79    /// Type classification of the entity.
80    pub entity_type: EntityType,
81    /// IDs of messages where this entity was mentioned.
82    pub message_ids: Vec<String>,
83    /// Total number of mentions across all messages.
84    pub mention_count: u32,
85    /// Computed importance score.
86    pub importance: f32,
87}
88
89/// An edge in the relationship graph.
90#[derive(Debug, Clone)]
91pub struct GraphEdge {
92    /// Source entity name.
93    pub from: String,
94    /// Target entity name.
95    pub to: String,
96    /// Type of relationship.
97    pub edge_type: EdgeType,
98    /// Strength of the relationship.
99    pub weight: f32,
100    /// Message ID where this edge was established, if any.
101    pub message_id: Option<String>,
102}
103
104/// Trait for querying an entity store.
105///
106/// Provides access to entity information without coupling to a specific storage
107/// implementation.
108pub trait EntityStoreT: Send + Sync {
109    /// Get entity names that match a given type.
110    fn entity_names_by_type(&self, entity_type: &EntityType) -> Vec<String>;
111
112    /// Get the top entities by mention count, returning (name, type) pairs.
113    fn top_entity_info(&self, limit: usize) -> Vec<(String, EntityType)>;
114}
115
116/// Trait for querying a relationship graph.
117///
118/// Provides read-only access to the graph structure without coupling to a
119/// specific implementation.
120pub trait RelationshipGraphT: Send + Sync {
121    /// Get a node by name.
122    fn get_node(&self, name: &str) -> Option<&GraphNode>;
123
124    /// Get all neighbor nodes.
125    fn get_neighbors(&self, name: &str) -> Vec<&GraphNode>;
126
127    /// Get all edges for a node.
128    fn get_edges(&self, name: &str) -> Vec<&GraphEdge>;
129
130    /// Search for nodes matching a query string.
131    fn search(&self, query: &str, limit: usize) -> Vec<&GraphNode>;
132
133    /// Find the shortest path between two nodes.
134    fn find_path(&self, from: &str, to: &str) -> Option<Vec<String>>;
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn test_entity_type_as_str() {
143        assert_eq!(EntityType::File.as_str(), "file");
144        assert_eq!(EntityType::Function.as_str(), "function");
145        assert_eq!(EntityType::Type.as_str(), "type");
146        assert_eq!(EntityType::Variable.as_str(), "variable");
147        assert_eq!(EntityType::Concept.as_str(), "concept");
148        assert_eq!(EntityType::Error.as_str(), "error");
149        assert_eq!(EntityType::Command.as_str(), "command");
150    }
151
152    #[test]
153    fn test_edge_type_weight() {
154        assert_eq!(EdgeType::Defines.weight(), 1.0);
155        assert_eq!(EdgeType::Contains.weight(), 0.9);
156        assert_eq!(EdgeType::DependsOn.weight(), 0.8);
157        assert_eq!(EdgeType::CoOccurs.weight(), 0.3);
158    }
159}