1use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
11pub enum EntityType {
12 File,
14 Function,
16 Type,
18 Variable,
20 Concept,
22 Error,
24 Command,
26}
27
28impl EntityType {
29 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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
45pub enum EdgeType {
46 CoOccurs,
48 Contains,
50 References,
52 DependsOn,
54 Modifies,
56 Defines,
58}
59
60impl EdgeType {
61 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#[derive(Debug, Clone)]
76pub struct GraphNode {
77 pub entity_name: String,
79 pub entity_type: EntityType,
81 pub message_ids: Vec<String>,
83 pub mention_count: u32,
85 pub importance: f32,
87}
88
89#[derive(Debug, Clone)]
91pub struct GraphEdge {
92 pub from: String,
94 pub to: String,
96 pub edge_type: EdgeType,
98 pub weight: f32,
100 pub message_id: Option<String>,
102}
103
104pub trait EntityStoreT: Send + Sync {
109 fn entity_names_by_type(&self, entity_type: &EntityType) -> Vec<String>;
111
112 fn top_entity_info(&self, limit: usize) -> Vec<(String, EntityType)>;
114}
115
116pub trait RelationshipGraphT: Send + Sync {
121 fn get_node(&self, name: &str) -> Option<&GraphNode>;
123
124 fn get_neighbors(&self, name: &str) -> Vec<&GraphNode>;
126
127 fn get_edges(&self, name: &str) -> Vec<&GraphEdge>;
129
130 fn search(&self, query: &str, limit: usize) -> Vec<&GraphNode>;
132
133 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}