use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum EntityType {
File,
Function,
Type,
Variable,
Concept,
Error,
Command,
}
impl EntityType {
pub fn as_str(&self) -> &'static str {
match self {
EntityType::File => "file",
EntityType::Function => "function",
EntityType::Type => "type",
EntityType::Variable => "variable",
EntityType::Concept => "concept",
EntityType::Error => "error",
EntityType::Command => "command",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum EdgeType {
CoOccurs,
Contains,
References,
DependsOn,
Modifies,
Defines,
}
impl EdgeType {
pub fn weight(&self) -> f32 {
match self {
EdgeType::Defines => 1.0,
EdgeType::Contains => 0.9,
EdgeType::DependsOn => 0.8,
EdgeType::Modifies => 0.7,
EdgeType::References => 0.6,
EdgeType::CoOccurs => 0.3,
}
}
}
#[derive(Debug, Clone)]
pub struct GraphNode {
pub entity_name: String,
pub entity_type: EntityType,
pub message_ids: Vec<String>,
pub mention_count: u32,
pub importance: f32,
}
#[derive(Debug, Clone)]
pub struct GraphEdge {
pub from: String,
pub to: String,
pub edge_type: EdgeType,
pub weight: f32,
pub message_id: Option<String>,
}
pub trait EntityStoreT: Send + Sync {
fn entity_names_by_type(&self, entity_type: &EntityType) -> Vec<String>;
fn top_entity_info(&self, limit: usize) -> Vec<(String, EntityType)>;
}
pub trait RelationshipGraphT: Send + Sync {
fn get_node(&self, name: &str) -> Option<&GraphNode>;
fn get_neighbors(&self, name: &str) -> Vec<&GraphNode>;
fn get_edges(&self, name: &str) -> Vec<&GraphEdge>;
fn search(&self, query: &str, limit: usize) -> Vec<&GraphNode>;
fn find_path(&self, from: &str, to: &str) -> Option<Vec<String>>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_entity_type_as_str() {
assert_eq!(EntityType::File.as_str(), "file");
assert_eq!(EntityType::Function.as_str(), "function");
assert_eq!(EntityType::Type.as_str(), "type");
assert_eq!(EntityType::Variable.as_str(), "variable");
assert_eq!(EntityType::Concept.as_str(), "concept");
assert_eq!(EntityType::Error.as_str(), "error");
assert_eq!(EntityType::Command.as_str(), "command");
}
#[test]
fn test_edge_type_weight() {
assert_eq!(EdgeType::Defines.weight(), 1.0);
assert_eq!(EdgeType::Contains.weight(), 0.9);
assert_eq!(EdgeType::DependsOn.weight(), 0.8);
assert_eq!(EdgeType::CoOccurs.weight(), 0.3);
}
}