use std::collections::HashMap;
use terraphim_types::{Concept, NormalizedTermValue, RoleName, Thesaurus};
#[derive(Debug, Clone)]
pub struct KnowledgeGraphRouter {
thesauri: HashMap<RoleName, Thesaurus>,
default_role: Option<RoleName>,
}
impl KnowledgeGraphRouter {
pub fn new() -> Self {
Self {
thesauri: HashMap::new(),
default_role: None,
}
}
pub fn with_default_role(mut self, role: RoleName) -> Self {
self.default_role = Some(role);
self
}
pub fn add_thesaurus(&mut self, role: RoleName, thesaurus: Thesaurus) {
self.thesauri.insert(role, thesaurus);
}
pub fn expand_terms(
&self,
terms: &[NormalizedTermValue],
role: Option<&RoleName>,
) -> Vec<NormalizedTermValue> {
let role = match role.or(self.default_role.as_ref()) {
Some(r) => r,
None => return terms.to_vec(),
};
let thesaurus = match self.thesauri.get(role) {
Some(t) => t,
None => return terms.to_vec(),
};
let mut expanded = terms.to_vec();
for term in terms {
if let Some(normalized) = thesaurus.get(term) {
let val = normalized.value.clone();
if !expanded.contains(&val) {
expanded.push(val);
}
}
}
expanded
}
pub fn score_provider_relevance(&self, _provider_id: &str, _concept: &Concept) -> f64 {
0.5
}
pub fn find_related_concepts(&self, _query: &str, _role: Option<&RoleName>) -> Vec<Concept> {
vec![]
}
}
impl Default for KnowledgeGraphRouter {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use terraphim_types::{NormalizedTerm, Thesaurus};
#[test]
fn test_kg_router_creation() {
let router = KnowledgeGraphRouter::new().with_default_role(RoleName::new("engineer"));
assert_eq!(router.default_role, Some(RoleName::new("engineer")));
}
#[test]
fn test_term_expansion() {
let mut router = KnowledgeGraphRouter::new();
let mut thesaurus = Thesaurus::new("programming".to_string());
let term = NormalizedTerm::new(1, NormalizedTermValue::from("rust"));
thesaurus.insert(NormalizedTermValue::from("rust"), term);
router.add_thesaurus(RoleName::new("engineer"), thesaurus);
let terms = vec![NormalizedTermValue::from("rust")];
let expanded = router.expand_terms(&terms, Some(&RoleName::new("engineer")));
assert!(!expanded.is_empty());
assert!(expanded.contains(&NormalizedTermValue::from("rust")));
}
}