ricecoder_domain_agents/
knowledge_base.rs

1//! Domain knowledge base management
2
3use crate::error::{DomainAgentError, Result};
4use crate::models::{Domain, KnowledgeBase, KnowledgeEntry};
5use ricecoder_storage::PathResolver;
6use tracing::{debug, info};
7
8/// Knowledge base manager for domain-specific agents
9pub struct KnowledgeBaseManager {
10    knowledge_bases: std::collections::HashMap<String, KnowledgeBase>,
11}
12
13impl KnowledgeBaseManager {
14    /// Create a new knowledge base manager
15    pub fn new() -> Self {
16        Self {
17            knowledge_bases: std::collections::HashMap::new(),
18        }
19    }
20}
21
22impl Default for KnowledgeBaseManager {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl KnowledgeBaseManager {
29
30    /// Load knowledge base for a domain
31    pub async fn load_knowledge_base(&mut self, domain: &str) -> Result<()> {
32        debug!("Loading knowledge base for domain: {}", domain);
33
34        let kb_path = self.get_knowledge_base_path(domain);
35
36        if kb_path.exists() {
37            let content = tokio::fs::read_to_string(&kb_path).await?;
38            let kb: KnowledgeBase = serde_yaml::from_str(&content)?;
39            self.knowledge_bases.insert(domain.to_string(), kb);
40            info!("Loaded knowledge base for domain: {}", domain);
41        } else {
42            debug!("Knowledge base not found for domain: {}, creating empty", domain);
43            let kb = KnowledgeBase::new(domain, "1.0.0");
44            self.knowledge_bases.insert(domain.to_string(), kb);
45        }
46
47        Ok(())
48    }
49
50    /// Get knowledge base for a domain
51    pub fn get_knowledge_base(&self, domain: &str) -> Result<&KnowledgeBase> {
52        self.knowledge_bases
53            .get(domain)
54            .ok_or_else(|| DomainAgentError::KnowledgeNotAvailable(domain.to_string()))
55    }
56
57    /// Get mutable knowledge base for a domain
58    pub fn get_knowledge_base_mut(&mut self, domain: &str) -> Result<&mut KnowledgeBase> {
59        self.knowledge_bases
60            .get_mut(domain)
61            .ok_or_else(|| DomainAgentError::KnowledgeNotAvailable(domain.to_string()))
62    }
63
64    /// Add knowledge entry to a domain
65    pub fn add_knowledge_entry(&mut self, domain: &str, entry: KnowledgeEntry) -> Result<()> {
66        let kb = self.get_knowledge_base_mut(domain)?;
67        kb.add_entry(entry);
68        Ok(())
69    }
70
71    /// Search knowledge by tag
72    pub fn search_by_tag(&self, domain: &str, tag: &str) -> Result<Vec<&KnowledgeEntry>> {
73        let kb = self.get_knowledge_base(domain)?;
74        Ok(kb.search_by_tag(tag))
75    }
76
77    /// Search knowledge by category
78    pub fn search_by_category(&self, domain: &str, category: &str) -> Result<Vec<&KnowledgeEntry>> {
79        let kb = self.get_knowledge_base(domain)?;
80        Ok(kb.search_by_category(category))
81    }
82
83    /// Save knowledge base to disk
84    pub async fn save_knowledge_base(&self, domain: &str) -> Result<()> {
85        let kb = self.get_knowledge_base(domain)?;
86        let kb_path = self.get_knowledge_base_path(domain);
87
88        // Create parent directory if it doesn't exist
89        if let Some(parent) = kb_path.parent() {
90            tokio::fs::create_dir_all(parent).await?;
91        }
92
93        let content = serde_yaml::to_string(kb)?;
94        tokio::fs::write(&kb_path, content).await?;
95        info!("Saved knowledge base for domain: {}", domain);
96
97        Ok(())
98    }
99
100    /// Get path to knowledge base file
101    fn get_knowledge_base_path(&self, domain: &str) -> std::path::PathBuf {
102        PathResolver::resolve_project_path()
103            .join("knowledge_bases")
104            .join(format!("{}.yaml", domain))
105    }
106
107    /// Load all knowledge bases for all domains
108    pub async fn load_all_knowledge_bases(&mut self) -> Result<()> {
109        debug!("Loading all knowledge bases");
110
111        for domain in Domain::all() {
112            self.load_knowledge_base(domain.as_str()).await?;
113        }
114
115        info!("Loaded all knowledge bases");
116        Ok(())
117    }
118
119    /// Get all loaded domains
120    pub fn get_loaded_domains(&self) -> Vec<&str> {
121        self.knowledge_bases.keys().map(|s| s.as_str()).collect()
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn test_knowledge_base_creation() {
131        let kb = KnowledgeBase::new("frontend", "1.0.0");
132        assert_eq!(kb.domain, "frontend");
133        assert_eq!(kb.version, "1.0.0");
134        assert!(kb.entries.is_empty());
135    }
136
137    #[test]
138    fn test_add_knowledge_entry() {
139        let mut kb = KnowledgeBase::new("frontend", "1.0.0");
140        let entry = KnowledgeEntry {
141            id: "react-hooks".to_string(),
142            category: "patterns".to_string(),
143            title: "React Hooks".to_string(),
144            description: "Best practices for using React Hooks".to_string(),
145            tags: vec!["react".to_string(), "hooks".to_string()],
146            example: Some("const [count, setCount] = useState(0);".to_string()),
147            references: vec!["https://react.dev/reference/react/hooks".to_string()],
148        };
149
150        kb.add_entry(entry);
151        assert_eq!(kb.entries.len(), 1);
152    }
153
154    #[test]
155    fn test_search_by_tag() {
156        let mut kb = KnowledgeBase::new("frontend", "1.0.0");
157        let entry = KnowledgeEntry {
158            id: "react-hooks".to_string(),
159            category: "patterns".to_string(),
160            title: "React Hooks".to_string(),
161            description: "Best practices for using React Hooks".to_string(),
162            tags: vec!["react".to_string(), "hooks".to_string()],
163            example: None,
164            references: vec![],
165        };
166
167        kb.add_entry(entry);
168        let results = kb.search_by_tag("react");
169        assert_eq!(results.len(), 1);
170    }
171
172    #[test]
173    fn test_search_by_category() {
174        let mut kb = KnowledgeBase::new("frontend", "1.0.0");
175        let entry = KnowledgeEntry {
176            id: "react-hooks".to_string(),
177            category: "patterns".to_string(),
178            title: "React Hooks".to_string(),
179            description: "Best practices for using React Hooks".to_string(),
180            tags: vec![],
181            example: None,
182            references: vec![],
183        };
184
185        kb.add_entry(entry);
186        let results = kb.search_by_category("patterns");
187        assert_eq!(results.len(), 1);
188    }
189}